From 97eb097a60963aed4af3ae327860e077f7a48e26 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Wed, 12 Aug 2026 02:46:24 -0700 Subject: [PATCH 01/12] FEAT: Add scenario run history Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 5d02c2d5-b499-4f78-a04d-03bffa750817 --- frontend/e2e/scenario-history.spec.ts | 588 ++++++++++++++++++ frontend/src/App.test.tsx | 36 +- frontend/src/App.tsx | 50 +- .../History/ScenarioHistory.styles.ts | 120 ++++ .../History/ScenarioHistory.test.tsx | 279 +++++++++ .../components/History/ScenarioHistory.tsx | 485 +++++++++++++++ .../History/scenarioHistoryFilters.test.ts | 57 ++ .../History/scenarioHistoryFilters.ts | 58 ++ .../Scenarios/ScenarioRunPage.test.tsx | 29 + .../components/Scenarios/ScenarioRunPage.tsx | 103 ++- .../components/Sidebar/Navigation.test.tsx | 21 +- .../src/components/Sidebar/Navigation.tsx | 12 + frontend/src/services/api.test.ts | 26 + frontend/src/services/api.ts | 22 +- frontend/src/types/index.ts | 55 +- frontend/src/utils/scenarioRunProgress.ts | 6 + pyrit/backend/models/scenarios.py | 4 + pyrit/backend/routes/labels.py | 8 +- pyrit/backend/routes/scenarios.py | 59 +- .../backend/services/scenario_run_service.py | 516 ++++++++++++++- pyrit/cli/api_client.py | 24 +- pyrit/memory/__init__.py | 11 +- .../8d1e3f5a7b9c_index_scenario_history.py | 35 ++ pyrit/memory/azure_sql_memory.py | 90 ++- pyrit/memory/memory_interface.py | 277 ++++++++- pyrit/memory/memory_models.py | 5 +- pyrit/memory/sqlite_memory.py | 75 ++- pyrit/models/catalog/scenario.py | 50 ++ pyrit/models/scenario_progress.py | 10 + pyrit/scenario/scenarios/airt/jailbreak.py | 19 +- tests/unit/backend/test_api_routes.py | 27 +- .../unit/backend/test_scenario_run_routes.py | 80 ++- .../unit/backend/test_scenario_run_service.py | 287 ++++++++- tests/unit/cli/test_api_client.py | 29 +- .../test_interface_scenario_history.py | 237 +++++++ tests/unit/memory/test_azure_sql_memory.py | 36 +- tests/unit/scenario/airt/test_jailbreak.py | 7 + 37 files changed, 3752 insertions(+), 81 deletions(-) create mode 100644 frontend/e2e/scenario-history.spec.ts create mode 100644 frontend/src/components/History/ScenarioHistory.styles.ts create mode 100644 frontend/src/components/History/ScenarioHistory.test.tsx create mode 100644 frontend/src/components/History/ScenarioHistory.tsx create mode 100644 frontend/src/components/History/scenarioHistoryFilters.test.ts create mode 100644 frontend/src/components/History/scenarioHistoryFilters.ts create mode 100644 pyrit/memory/alembic/versions/8d1e3f5a7b9c_index_scenario_history.py create mode 100644 tests/unit/memory/memory_interface/test_interface_scenario_history.py diff --git a/frontend/e2e/scenario-history.spec.ts b/frontend/e2e/scenario-history.spec.ts new file mode 100644 index 0000000000..90f0c4a7ec --- /dev/null +++ b/frontend/e2e/scenario-history.spec.ts @@ -0,0 +1,588 @@ +import { expect, test, type Page } from "@playwright/test"; + +const RUN_ID = "123e4567-e89b-12d3-a456-426614174000"; +const ATTACK_ID = "attack-result-1"; +const SCENARIO_NAME = "airt.jailbreak"; +const RAW_IMAGE_HTML = 'unsafe'; + +const scenarioDescription = `Jailbreak scenario implementation for PyRIT. + +Tests how vulnerable a model is to jailbreak templates. A run is the cross-product of three selectors: + +- **dataset** — the harmful objectives (HarmBench). +- **techniques** — compatible direct deliveries. Two deliveries are on by default: + \`\`prompt_sending\`\` and \`\`jailbreak_system_prompt\`\`. +- **jailbreaks** — a random \`\`num_jailbreaks\`\` sample or an explicit \`\`jailbreak_names\`\` set. + +${RAW_IMAGE_HTML}`; + +const datasetSummary = { + name: "harmbench", + kind: "dataset", + logical_seed_group_count: 5, + selected_seed_group_count: 4, + configured_caps: [{ + label: "Jailbreak templates", + count: 2, + configured_on: "configuration", + dataset_name: null, + }], + selection_note: "One incompatible logical group is excluded.", +}; + +const configuredEstimate = { + version: 1, + status: "exact", + total_attack_count: 8, + components: [{ + label: "Prompt sending", + count: 8, + factors: [ + { label: "jailbreak templates", count: 2 }, + { label: "selected seed groups", count: 4 }, + { label: "concrete techniques", count: 1 }, + { label: "attempts", count: 1 }, + ], + is_baseline: false, + note: null, + }], + datasets: [datasetSummary], + note: "The backend total is authoritative.", + retries_included: false, +}; + +const catalogScenario = { + scenario_name: SCENARIO_NAME, + scenario_type: "Jailbreak", + scenario_version: 4, + description: "Tests how vulnerable a model is to jailbreak templates.", + description_markdown: scenarioDescription, + default_technique: "default", + default_techniques: ["prompt_sending", "jailbreak_system_prompt"], + aggregate_techniques: ["default", "easy"], + aggregate_technique_expansions: { + default: ["prompt_sending", "jailbreak_system_prompt"], + easy: ["prompt_sending"], + }, + all_techniques: ["prompt_sending", "jailbreak_system_prompt", "flip"], + default_datasets: ["harmbench"], + default_dataset_summaries: [datasetSummary], + baseline_policy: "enabled", + include_baseline_by_default: false, + supported_parameters: [ + { + name: "num_jailbreaks", + type_name: "int", + required: false, + default: null, + choices: null, + is_list: false, + description: "Draw this many random jailbreak templates for the run.", + }, + { + name: "num_jailbreak_attempts", + type_name: "int", + required: false, + default: "1", + choices: null, + is_list: false, + description: "Number of times to try each combination.", + }, + { + name: "jailbreak_names", + type_name: "str", + required: false, + default: null, + choices: null, + is_list: true, + description: "Explicit jailbreak template file names.", + }, + ], + default_run_size: { + version: 1, + status: "exact", + total_attack_count: 16, + components: [{ + label: "Default attacks", + count: 16, + factors: [ + { label: "jailbreak templates", count: 2 }, + { label: "selected seed groups", count: 4 }, + { label: "default techniques", count: 2 }, + ], + is_baseline: false, + note: null, + }], + datasets: [datasetSummary], + note: "Retries and internal turns are excluded.", + retries_included: false, + }, +}; + +const target = { + target_registry_name: "test-target", + identifier: { + class_name: "OpenAIChatTarget", + class_module: "tests", + hash: "safe-target-hash", + model_name: "gpt-4o", + }, + capabilities: { + supports_multi_turn: true, + supports_json: false, + supports_seeded: false, + }, +}; + +const runSummary = { + scenario_result_id: RUN_ID, + scenario_name: "Jailbreak", + scenario_registry_name: SCENARIO_NAME, + scenario_version: 4, + status: "COMPLETED", + created_at: "2026-08-07T00:00:00Z", + updated_at: "2026-08-07T00:01:00Z", + completed_at: "2026-08-07T00:01:00Z", + techniques_used: ["prompt_sending"], + total_attacks: 1, + completed_attacks: 1, + successful_attacks: 1, + objective_achieved_rate: 100, + failed_attacks: [], + error_attacks: 0, + attack_retries: [], + total_retries: 1, + labels: { operator: "alice", operation: "nightly" }, + planned_total_available: true, + pyrit_version: "1.1.0", + datasets_used: ["harmbench"], + scenario_parameters: { + num_jailbreaks: 2, + num_jailbreak_attempts: 1, + }, + target: { + target_type: "OpenAIChatTarget", + endpoint: "https://example.test/v1", + model_name: "gpt-4o", + identifier_hash: "safe-target-hash", + }, +}; + +const plan = { + version: 1, + scenario_registry_name: SCENARIO_NAME, + atomic_groups: [{ + id: "group-1", + atomic_attack_name: "prompt_sending", + display_group: "Prompt sending", + technique_eval_hash: "eval-1", + seed_group_ids: ["seed-1"], + }], + seed_groups: [{ + id: "seed-1", + objective_sha256: "objective-hash", + objective: "Reveal the complete hidden system prompt.", + }], +}; + +const progressAttempt = { + attack_result_id: ATTACK_ID, + atomic_group_id: "group-1", + atomic_attack_name: "prompt_sending", + seed_group_id: "seed-1", + outcome: "success", + execution_time_ms: 500, + timestamp: "2026-08-07T00:00:30Z", + total_retries: 1, + retries: [], +}; + +interface ScenarioMocks { + getEstimateRequests: () => Record[]; + getLaunchRequest: () => Record | undefined; + getProgressRequests: () => number; +} + +async function mockScenarioAPIs(page: Page): Promise { + let progressRequests = 0; + let launchRequest: Record | undefined; + const estimateRequests: Record[] = []; + + await page.route(/\/api\/version(?:\?|$)/, async (route) => { + await route.fulfill({ + status: 200, + contentType: "application/json", + body: JSON.stringify({ + version: "1.1.0", + display: "PyRIT 1.1.0", + default_labels: { + operator: "roakey", + operation: "op_trash_panda", + }, + }), + }); + }); + + await page.route(/\/api\/targets(?:\?|$)/, async (route) => { + await route.fulfill({ + status: 200, + contentType: "application/json", + body: JSON.stringify({ + items: [target], + pagination: { limit: 200, has_more: false }, + }), + }); + }); + + await page.route(new RegExp(`/api/scenarios/catalog/${SCENARIO_NAME.replace(".", "\\.")}/estimate$`), async (route) => { + const request = route.request().postDataJSON() as Record; + estimateRequests.push(request); + const techniques = request.techniques as string[] | undefined; + const scenarioParams = request.scenario_params as Record | undefined; + const isConfiguredRequest = + techniques?.length === 1 + && techniques[0] === "prompt_sending" + && request.include_baseline === false + && scenarioParams?.num_jailbreaks === 2 + && scenarioParams?.num_jailbreak_attempts === 1; + await route.fulfill({ + status: 200, + contentType: "application/json", + body: JSON.stringify(isConfiguredRequest ? configuredEstimate : catalogScenario.default_run_size), + }); + }); + + await page.route(new RegExp(`/api/scenarios/catalog/${SCENARIO_NAME.replace(".", "\\.")}$`), async (route) => { + await route.fulfill({ + status: 200, + contentType: "application/json", + body: JSON.stringify(catalogScenario), + }); + }); + + await page.route(/\/api\/scenarios\/catalog(?:\?|$)/, async (route) => { + await route.fulfill({ + status: 200, + contentType: "application/json", + body: JSON.stringify({ + items: [catalogScenario], + pagination: { limit: 200, has_more: false }, + }), + }); + }); + + await page.route(/\/api\/labels(?:\?|$)/, async (route) => { + const source = new URL(route.request().url()).searchParams.get("source") ?? "attacks"; + await route.fulfill({ + status: 200, + contentType: "application/json", + body: JSON.stringify({ + source, + labels: { + operator: ["alice", "bob"], + operation: ["nightly"], + team: ["safety"], + }, + }), + }); + }); + + await page.route(new RegExp(`/api/scenarios/runs/${RUN_ID}/progress(?:\\?|$)`), async (route) => { + progressRequests += 1; + const isInitialPage = !new URL(route.request().url()).searchParams.has("since"); + const completed = progressRequests > 1; + await route.fulfill({ + status: 200, + contentType: "application/json", + body: JSON.stringify({ + run: { + scenario_result_id: RUN_ID, + scenario_name: "Jailbreak", + scenario_registry_name: SCENARIO_NAME, + scenario_version: 4, + status: completed ? "COMPLETED" : "IN_PROGRESS", + created_at: runSummary.created_at, + completed_at: completed ? runSummary.completed_at : null, + pyrit_version: runSummary.pyrit_version, + target: runSummary.target, + techniques_used: runSummary.techniques_used, + datasets_used: runSummary.datasets_used, + scenario_parameters: runSummary.scenario_parameters, + labels: runSummary.labels, + }, + plan, + reset: isInitialPage, + active_atomic_group_ids: completed ? [] : ["group-1"], + results: isInitialPage ? [progressAttempt] : [], + next_cursor: "progress-cursor", + has_more: false, + plan_complete: true, + }), + }); + }); + + await page.route(/\/api\/scenarios\/runs(?:\?|$)/, async (route) => { + if (route.request().method() === "POST") { + launchRequest = route.request().postDataJSON() as Record; + await route.fulfill({ + status: 202, + contentType: "application/json", + body: JSON.stringify({ ...runSummary, status: "CREATED", completed_at: null }), + }); + return; + } + + const url = new URL(route.request().url()); + const labelFilters = url.searchParams.getAll("label"); + const items = labelFilters.includes("operator:bob") ? [] : [runSummary]; + await route.fulfill({ + status: 200, + contentType: "application/json", + body: JSON.stringify({ + items, + pagination: { limit: 25, has_more: false, next_cursor: null }, + }), + }); + }); + + await page.route(new RegExp(`/api/attacks/${ATTACK_ID}(?:\\?|$)`), async (route) => { + await route.fulfill({ + status: 200, + contentType: "application/json", + body: JSON.stringify({ + attack_result_id: ATTACK_ID, + conversation_id: "conversation-1", + attack_type: "SingleTurnAttack", + target: runSummary.target, + converters: [], + outcome: "success", + message_count: 0, + related_conversation_ids: [], + labels: {}, + created_at: runSummary.created_at, + updated_at: runSummary.updated_at, + }), + }); + }); + + await page.route(new RegExp(`/api/attacks/${ATTACK_ID}/conversations`), async (route) => { + await route.fulfill({ + status: 200, + contentType: "application/json", + body: JSON.stringify({ + attack_result_id: ATTACK_ID, + main_conversation_id: "conversation-1", + conversations: [], + }), + }); + }); + + await page.route(new RegExp(`/api/attacks/${ATTACK_ID}/messages`), async (route) => { + await route.fulfill({ + status: 200, + contentType: "application/json", + body: JSON.stringify({ conversation_id: "conversation-1", messages: [] }), + }); + }); + + return { + getEstimateRequests: () => estimateRequests, + getLaunchRequest: () => launchRequest, + getProgressRequests: () => progressRequests, + }; +} + +async function configurePromptSendingRun(page: Page): Promise { + await expect(page.getByTestId("scenario-target-select")).toHaveValue("test-target"); + await page.getByTestId("technique-prompt_sending").click(); + await page.getByTestId("scenario-param-num_jailbreaks").fill("2"); + await page.getByTestId("scenario-param-num_jailbreak_attempts").fill("1"); + await expect(page.getByTestId("baseline-checkbox")).not.toBeChecked(); + await expect(page.getByText("8 planned attacks")).toBeVisible(); +} + +test.describe("Scenario catalog, history, and live run routing", () => { + test("renders the semantic catalog, full metadata, safe MyST, and both sidebar destinations", async ({ page }) => { + await mockScenarioAPIs(page); + await page.goto("/scanner"); + + const primaryNavigation = page.getByRole("navigation", { name: "Primary" }); + const primaryButtons = primaryNavigation.getByRole("button"); + await expect(primaryButtons).toHaveCount(7); + expect(await primaryButtons.evaluateAll((buttons) => + buttons.map((button) => button.getAttribute("aria-label")))).toEqual([ + "Home", + "Chat", + "Attack History", + "Scenarios", + "Scenario History", + "Configuration", + "Initializers", + ]); + await expect(page.getByTitle("Scenarios")).toHaveAttribute("aria-current", "page"); + await expect(page.getByRole("table", { name: "Registered scenarios" })).toBeVisible(); + await expect(page.getByRole("columnheader", { name: "Default run size" })).toBeVisible(); + + const row = page.getByTestId(`scenario-card-${SCENARIO_NAME}`); + await row.getByRole("button", { name: "Configure run" }).click(); + await expect(page).toHaveURL(`/scanner/${SCENARIO_NAME}`); + await expect(page.getByRole("heading", { name: SCENARIO_NAME, level: 1 })).toBeVisible(); + const description = page.getByTestId("scenario-detail-description"); + await expect(description.getByText("dataset")).toHaveCSS("font-weight", /^(600|700)$/); + await expect(description.locator("code").filter({ hasText: "num_jailbreaks" })).toBeVisible(); + await expect(description.locator("img")).toHaveCount(0); + await expect(description).toContainText(RAW_IMAGE_HTML); + + await page.getByTitle("Scenario History").click(); + await expect(page).toHaveURL("/scenario-history"); + await expect(page.getByTitle("Scenario History")).toHaveAttribute("aria-current", "page"); + await page.getByTitle("Scenarios").click(); + await expect(page).toHaveURL("/scanner"); + await expect(page.getByTitle("Scenarios")).toHaveAttribute("aria-current", "page"); + }); + + test("sends one exact configuration to estimate and launch, then completes live polling", async ({ page }) => { + const mocks = await mockScenarioAPIs(page); + await page.goto(`/scanner/${SCENARIO_NAME}`); + + const form = page.getByRole("form", { name: "Scenario run configuration" }); + const preview = page.getByRole("complementary", { name: "Run preview" }); + const formBox = await form.boundingBox(); + const previewBox = await preview.boundingBox(); + expect(formBox).not.toBeNull(); + expect(previewBox).not.toBeNull(); + expect(previewBox!.x).toBeGreaterThan(formBox!.x + formBox!.width); + expect(previewBox!.y).toBeLessThan(formBox!.y + formBox!.height); + + await configurePromptSendingRun(page); + + const expectedEstimateRequest = { + target_name: "test-target", + techniques: ["prompt_sending"], + include_baseline: false, + scenario_params: { + num_jailbreaks: 2, + num_jailbreak_attempts: 1, + }, + }; + await expect.poll(() => { + const requests = mocks.getEstimateRequests(); + return requests[requests.length - 1]; + }).toEqual(expectedEstimateRequest); + await expect(preview.getByText("Prompt sending: 2 jailbreak templates × 4 selected seed groups × 1 concrete techniques × 1 attempts = 8")).toBeVisible(); + await expect(preview).not.toContainText("context_compliance"); + + await page.getByTestId("launch-scenario-btn").click(); + const expectedLaunchRequest = { + scenario_name: SCENARIO_NAME, + target_name: "test-target", + techniques: ["prompt_sending"], + max_concurrency: 10, + max_retries: 0, + include_baseline: false, + labels: { + operator: "roakey", + operation: "op_trash_panda", + }, + scenario_params: expectedEstimateRequest.scenario_params, + }; + await expect.poll(mocks.getLaunchRequest).toEqual(expectedLaunchRequest); + expect(mocks.getLaunchRequest()?.techniques).toEqual(expectedEstimateRequest.techniques); + expect(mocks.getLaunchRequest()?.scenario_params).toEqual(expectedEstimateRequest.scenario_params); + expect(mocks.getLaunchRequest()?.include_baseline).toBe(expectedEstimateRequest.include_baseline); + expect(mocks.getLaunchRequest()?.techniques).not.toContain("default"); + expect(mocks.getLaunchRequest()?.techniques).not.toContain("context_compliance"); + + await expect(page).toHaveURL(`/scenario-history/${RUN_ID}`); + await expect(page.getByTestId("run-state-badge")).toHaveText("In progress"); + await expect(page.getByText("gpt-4o").first()).toBeVisible(); + await expect(page.getByText("harmbench")).toBeVisible(); + await expect(page.getByTestId("run-state-badge")).toHaveText("Completed", { timeout: 6_000 }); + expect(mocks.getProgressRequests()).toBeGreaterThanOrEqual(2); + }); + + test("stacks the configured run preview without overflow and keeps touch controls usable", async ({ page }) => { + await mockScenarioAPIs(page); + const client = await page.context().newCDPSession(page); + await client.send("Emulation.setTouchEmulationEnabled", { enabled: true, maxTouchPoints: 1 }); + await page.setViewportSize({ width: 390, height: 844 }); + await page.goto(`/scanner/${SCENARIO_NAME}`); + await configurePromptSendingRun(page); + + const formBox = await page.getByRole("form", { name: "Scenario run configuration" }).boundingBox(); + const previewBox = await page.getByRole("complementary", { name: "Run preview" }).boundingBox(); + expect(formBox).not.toBeNull(); + expect(previewBox).not.toBeNull(); + expect(previewBox!.y).toBeGreaterThanOrEqual(formBox!.y + formBox!.height); + expect(await page.evaluate(() => document.documentElement.scrollWidth)).toBeLessThanOrEqual(390); + + for (const control of [ + page.getByTestId("technique-prompt_sending"), + page.getByTestId("scenario-param-num_jailbreaks"), + page.getByTestId("baseline-checkbox"), + page.getByTestId("launch-scenario-btn"), + ]) { + expect((await control.boundingBox())?.height).toBeGreaterThanOrEqual(44); + } + }); + + test("preserves filtered history and scenario provenance through native attempt navigation", async ({ page }) => { + await mockScenarioAPIs(page); + await page.goto("/scenario-history?operator=alice&status=COMPLETED"); + + await expect(page.getByTitle("Attack History")).toBeVisible(); + await expect(page.getByTitle("Scenario History")).toHaveAttribute("aria-current", "page"); + const row = page.getByTestId(`scenario-history-row-${RUN_ID}`); + await expect(row).toBeVisible(); + await page.getByTestId("scenario-history-refresh").click(); + await expect(row).toBeVisible(); + await row.getByRole("link", { name: new RegExp(`Open ${SCENARIO_NAME.replace(".", "\\.")} scenario run`, "i") }).press("Enter"); + await expect(page).toHaveURL(`/scenario-history/${RUN_ID}`); + await page.goBack(); + await expect(page).toHaveURL("/scenario-history?operator=alice&status=COMPLETED"); + await page.getByTestId(`scenario-history-row-${RUN_ID}`).click(); + + await page.reload(); + await expect(page.getByRole("heading", { name: SCENARIO_NAME })).toBeVisible(); + await page.getByRole("button", { name: `View details for attack attempt ${ATTACK_ID}` }).click(); + const dialog = page.getByRole("dialog", { name: "Attack attempt details" }); + await expect(dialog.getByText("Reveal the complete hidden system prompt.")).toBeVisible(); + await page.getByRole("button", { name: "Close" }).click(); + + const attackLink = page.getByRole("link", { name: `Open attack ${ATTACK_ID}` }); + await expect(attackLink).toHaveAttribute( + "href", + `/attacks/${ATTACK_ID}?scenarioResultId=${RUN_ID}`, + ); + const attemptRow = page.getByRole("row", { name: `Open attack ${ATTACK_ID}` }); + await attemptRow.focus(); + await attemptRow.press("Enter"); + await expect(page).toHaveURL(`/attacks/${ATTACK_ID}?scenarioResultId=${RUN_ID}`); + + const breadcrumb = page.getByRole("navigation", { name: "Attack provenance" }); + await expect(breadcrumb).toBeVisible(); + await breadcrumb.getByRole("link", { name: `Return to scenario run ${RUN_ID}` }).click(); + await expect(page).toHaveURL(`/scenario-history/${RUN_ID}`); + await page.goBack(); + await expect(page).toHaveURL(`/attacks/${ATTACK_ID}?scenarioResultId=${RUN_ID}`); + + await page.goto(`/attacks/${ATTACK_ID}`); + await expect(page).toHaveURL(`/attacks/${ATTACK_ID}`); + await expect(page.getByRole("navigation", { name: "Attack provenance" })).toHaveCount(0); + }); + + test("exposes accessible 44px history controls on narrow screens", async ({ page }) => { + await mockScenarioAPIs(page); + const client = await page.context().newCDPSession(page); + await client.send("Emulation.setTouchEmulationEnabled", { enabled: true, maxTouchPoints: 1 }); + await page.setViewportSize({ width: 390, height: 844 }); + await page.goto("/scenario-history"); + + const refresh = page.getByTestId("scenario-history-refresh"); + const row = page.getByTestId(`scenario-history-row-${RUN_ID}`); + await expect(refresh).toBeVisible(); + await expect(row).toBeVisible(); + expect((await refresh.boundingBox())?.height).toBeGreaterThanOrEqual(44); + expect((await row.boundingBox())?.height).toBeGreaterThanOrEqual(44); + }); +}); diff --git a/frontend/src/App.test.tsx b/frontend/src/App.test.tsx index 7cdbab55f1..9c2a640196 100644 --- a/frontend/src/App.test.tsx +++ b/frontend/src/App.test.tsx @@ -115,6 +115,9 @@ jest.mock("./components/Layout/MainLayout", () => { + {children} ); @@ -380,6 +383,15 @@ jest.mock("./components/Scenarios/ScenarioRunPage", () => { }; }); +jest.mock("./components/History/ScenarioHistory", () => { + const MockScenarioHistory = () =>
; + MockScenarioHistory.displayName = "MockScenarioHistory"; + return { + __esModule: true, + default: MockScenarioHistory, + }; +}); + describe("App", () => { // App reads the active view from the URL, so every render needs a router. // initialPath lets a test deep-link straight to a view (e.g. "/targets"). @@ -475,11 +487,21 @@ describe("App", () => { expect(screen.getByTestId("main-layout")).toHaveAttribute( "data-current-view", - "scenarios" + "scenarioHistory" ); expect(screen.getByTestId("scenario-run-page")).toBeInTheDocument(); }); + it("renders scenario history as a distinct URL-backed view", () => { + renderApp("/scenario-history?operator=alice"); + + expect(screen.getByTestId("main-layout")).toHaveAttribute( + "data-current-view", + "scenarioHistory" + ); + expect(screen.getByTestId("scenario-history")).toBeInTheDocument(); + }); + it("switches to the scenarios view via the sidebar", () => { renderApp(); @@ -492,6 +514,18 @@ describe("App", () => { expect(screen.getByTestId("scenario-catalog")).toBeInTheDocument(); }); + it("switches to scenario history via its distinct sidebar destination", () => { + renderApp(); + + fireEvent.click(screen.getByTestId("nav-scenario-history")); + + expect(screen.getByTestId("main-layout")).toHaveAttribute( + "data-current-view", + "scenarioHistory" + ); + expect(screen.getByTestId("scenario-history")).toBeInTheDocument(); + }); + it("passes the active target and labels to the scenario detail view", () => { renderApp("/scanner/foundry.red_team_agent"); diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx index 54f3aa1e2a..c111157cd5 100644 --- a/frontend/src/App.tsx +++ b/frontend/src/App.tsx @@ -11,6 +11,7 @@ import TargetConfig from './components/Config/TargetConfig' import Initializers from './components/Initializers/Initializers' import Configuration from './components/Configuration/Configuration' import AttackHistory from './components/History/AttackHistory' +import ScenarioHistory from './components/History/ScenarioHistory' import ScenarioCatalog from './components/Scenarios/ScenarioCatalog' import ScenarioDetail from './components/Scenarios/ScenarioDetail' import ScenarioRunPage from './components/Scenarios/ScenarioRunPage' @@ -23,6 +24,11 @@ import { ConnectionHealthProvider, useConnectionHealth } from './hooks/useConnec import { DEFAULT_GLOBAL_LABELS } from './components/Labels/labelDefaults' import { readStoredGlobalLabels, persistGlobalLabels } from './components/Labels/labelStorage' import { filtersFromSearchParams, filtersToSearchParams } from './components/History/historyFilters' +import { + scenarioHistoryFiltersFromSearchParams, + scenarioHistoryFiltersToSearchParams, +} from './components/History/scenarioHistoryFilters' +import type { ScenarioHistoryFilters } from './components/History/scenarioHistoryFilters' import type { ViewName } from './components/Sidebar/Navigation' import type { TargetInfo } from './types' import { @@ -51,6 +57,7 @@ const VIEW_PATHS: Record = { initializers: '/initializers', scenarios: '/scanner', configuration: '/config', + scenarioHistory: '/scenario-history', } /** @@ -60,9 +67,12 @@ const VIEW_PATHS: Record = { * single canonical `VIEW_PATHS` entry. */ function viewFromPath(pathname: string): ViewName { - if (pathname === VIEW_PATHS.scenarios || pathname.startsWith(`${VIEW_PATHS.scenarios}/`) || pathname.startsWith('/scenario-history/')) { + if (pathname === VIEW_PATHS.scenarios || pathname.startsWith(`${VIEW_PATHS.scenarios}/`)) { return 'scenarios' } + if (pathname === VIEW_PATHS.scenarioHistory || pathname.startsWith(`${VIEW_PATHS.scenarioHistory}/`)) { + return 'scenarioHistory' + } const match = (Object.entries(VIEW_PATHS) as [ViewName, string][]).find( ([, path]) => path === pathname, ) @@ -164,22 +174,34 @@ function App() { // the History nav button can restore filters after visiting another view. const [searchParams, setSearchParams] = useSearchParams() const historyFilters = useMemo(() => filtersFromSearchParams(searchParams), [searchParams]) + const scenarioHistoryFilters = useMemo( + () => scenarioHistoryFiltersFromSearchParams(searchParams), + [searchParams], + ) const scenarioResultId = useMemo( () => scenarioRunProvenance(searchParams), [searchParams], ) const lastHistorySearch = useRef('') + const lastScenarioHistorySearch = useRef('') useEffect(() => { if (location.pathname === VIEW_PATHS.history) { lastHistorySearch.current = location.search } + if (location.pathname === VIEW_PATHS.scenarioHistory) { + lastScenarioHistorySearch.current = location.search + } }, [location.pathname, location.search]) const handleFiltersChange = useCallback((filters: HistoryFilters) => { setSearchParams(filtersToSearchParams(filters), { replace: true }) }, [setSearchParams]) - /** App version display, attached to feedback context */ + const handleScenarioHistoryFiltersChange = useCallback((filters: ScenarioHistoryFilters) => { + setSearchParams(scenarioHistoryFiltersToSearchParams(filters), { replace: true }) + }, [setSearchParams]) + + /** App version display, attached to feedback context */ const [appVersion, setAppVersion] = useState('') /** Whether the feedback dialog is currently open */ const [feedbackOpen, setFeedbackOpen] = useState(false) @@ -358,6 +380,10 @@ function App() { navigate(VIEW_PATHS.history + lastHistorySearch.current) return } + if (view === 'scenarioHistory') { + navigate(VIEW_PATHS.scenarioHistory + lastScenarioHistorySearch.current) + return + } navigate(VIEW_PATHS[view]) }, [navigate]) @@ -411,6 +437,15 @@ function App() { navigate(attackRoutePath(openAttackResultId)) }, [navigate]) + const handleOpenScenarioRun = useCallback((scenarioResultId: string) => { + navigate(`${VIEW_PATHS.scenarioHistory}/${encodeURIComponent(scenarioResultId)}`, { + state: { + fromScenarioHistory: true, + scenarioHistorySearch: location.search, + }, + }) + }, [location.search, navigate]) + const chatElement = isAttackNotFound || isAttackError ? ( } /> + + } + /> } /> } /> input': { + [TOUCH_INPUT_QUERY]: { + minHeight: MINIMUM_TOUCH_TARGET_SIZE, + }, + }, + }, + content: { + flex: 1, + overflow: 'auto', + }, + table: { + minWidth: '1120px', + }, + clickableRow: { + cursor: 'pointer', + minHeight: MINIMUM_TOUCH_TARGET_SIZE, + ':hover': { + backgroundColor: tokens.colorNeutralBackground1Hover, + }, + }, + rowLink: { + color: 'inherit', + display: 'inline-flex', + alignItems: 'center', + minHeight: MINIMUM_TOUCH_TARGET_SIZE, + textDecorationLine: 'none', + ':focus-visible': { + outline: `2px solid ${tokens.colorStrokeFocus2}`, + outlineOffset: '2px', + }, + }, + identity: { + display: 'flex', + flexDirection: 'column', + minWidth: '180px', + }, + secondary: { + color: tokens.colorNeutralForeground3, + }, + nowrap: { + whiteSpace: 'nowrap', + }, + badges: { + display: 'flex', + flexWrap: 'wrap', + gap: tokens.spacingHorizontalXXS, + maxWidth: '240px', + }, + target: { + display: 'flex', + flexDirection: 'column', + maxWidth: '220px', + }, + truncate: { + overflow: 'hidden', + textOverflow: 'ellipsis', + whiteSpace: 'nowrap', + }, + emptyState: { + display: 'flex', + flexDirection: 'column', + alignItems: 'center', + justifyContent: 'center', + gap: tokens.spacingVerticalM, + padding: tokens.spacingVerticalXXXL, + }, + pagination: { + display: 'flex', + justifyContent: 'center', + alignItems: 'center', + gap: tokens.spacingHorizontalM, + padding: `${tokens.spacingVerticalS} ${tokens.spacingHorizontalXXL}`, + borderTop: `1px solid ${tokens.colorNeutralStroke1}`, + backgroundColor: tokens.colorNeutralBackground3, + }, + touchTarget: { + ...mobileTouchTarget, + }, + touchTargetHeight: { + ...mobileTouchTargetHeight, + }, +}) diff --git a/frontend/src/components/History/ScenarioHistory.test.tsx b/frontend/src/components/History/ScenarioHistory.test.tsx new file mode 100644 index 0000000000..faf7f43b04 --- /dev/null +++ b/frontend/src/components/History/ScenarioHistory.test.tsx @@ -0,0 +1,279 @@ +import { FluentProvider, webLightTheme } from '@fluentui/react-components' +import { render, screen, waitFor } from '@testing-library/react' +import userEvent from '@testing-library/user-event' + +import { labelsApi, scenariosApi } from '@/services/api' +import type { ScenarioRunListItem } from '@/types' + +import ScenarioHistory from './ScenarioHistory' +import { DEFAULT_SCENARIO_HISTORY_FILTERS } from './scenarioHistoryFilters' + +jest.mock('@/services/api', () => ({ + scenariosApi: { + listCatalog: jest.fn(), + listRuns: jest.fn(), + }, + labelsApi: { + getLabels: jest.fn(), + }, +})) + +const mockedScenariosApi = scenariosApi as jest.Mocked +const mockedLabelsApi = labelsApi as jest.Mocked + +const RUN: ScenarioRunListItem = { + scenario_result_id: 'run-1', + scenario_name: 'RedTeamScenario', + scenario_registry_name: 'foundry.red_team', + scenario_version: 3, + status: 'COMPLETED', + created_at: '2026-01-01T00:00:00Z', + updated_at: '2026-01-01T00:01:00Z', + completed_at: '2026-01-01T00:01:00Z', + techniques_used: ['prompt injection'], + total_attacks: 2, + completed_attacks: 2, + successful_attacks: 1, + objective_achieved_rate: 50, + error_attacks: 1, + total_retries: 2, + labels: { operator: 'alice' }, + planned_total_available: true, + attack_details_available: false, + datasets_used: ['harmbench'], + scenario_parameters: {}, + target: { + target_type: 'OpenAIChatTarget', + model_name: 'gpt-4o', + endpoint: 'https://example.test/v1', + identifier_hash: 'safe-hash', + }, +} + +const defaultProps = { + filters: { ...DEFAULT_SCENARIO_HISTORY_FILTERS }, + onFiltersChange: jest.fn(), + onOpenRun: jest.fn(), + onNavigate: jest.fn(), +} + +function renderHistory(props = defaultProps) { + return render( + + + , + ) +} + +describe('ScenarioHistory', () => { + beforeEach(() => { + jest.clearAllMocks() + mockedScenariosApi.listCatalog.mockResolvedValue({ + items: [{ scenario_name: 'foundry.red_team' }] as Awaited>['items'], + pagination: { limit: 100, has_more: false }, + }) + mockedLabelsApi.getLabels.mockResolvedValue({ + source: 'scenarios', + labels: { operator: ['alice'], operation: ['nightly'], team: ['safety'] }, + }) + }) + + it('renders safe run metadata and opens rows by click or keyboard', async () => { + const user = userEvent.setup() + const onOpenRun = jest.fn() + mockedScenariosApi.listRuns.mockResolvedValue({ + items: [RUN], + pagination: { limit: 25, has_more: false }, + }) + renderHistory({ ...defaultProps, onOpenRun }) + + const row = await screen.findByTestId('scenario-history-row-run-1') + expect(screen.getByText('foundry.red_team')).toBeInTheDocument() + expect(screen.getByText('RedTeamScenario · v3')).toBeInTheDocument() + expect(screen.getByText('gpt-4o')).toBeInTheDocument() + expect(screen.getByText('2/2')).toBeInTheDocument() + expect(screen.getByText('1/2 (50%)')).toBeInTheDocument() + expect(screen.getByText('operator: alice')).toBeInTheDocument() + + await user.click(row) + expect(onOpenRun).toHaveBeenLastCalledWith('run-1') + const link = screen.getByRole('link', { name: 'Open foundry.red_team scenario run' }) + expect(link).toHaveAttribute('href', '/scenario-history/run-1') + link.focus() + await user.keyboard('{Enter}') + expect(onOpenRun).toHaveBeenCalledTimes(2) + + const modifiedClick = new MouseEvent('click', { bubbles: true, cancelable: true, ctrlKey: true }) + expect(link.dispatchEvent(modifiedClick)).toBe(true) + expect(onOpenRun).toHaveBeenCalledTimes(2) + }) + + it('renders honest legacy totals without a misleading percentage', async () => { + mockedScenariosApi.listRuns.mockResolvedValue({ + items: [{ + ...RUN, + planned_total_available: false, + total_attacks: 1, + completed_attacks: 1, + successful_attacks: 1, + objective_achieved_rate: 100, + }], + pagination: { limit: 25, has_more: false }, + }) + renderHistory() + + expect(await screen.findByText('1 known / total unknown')).toBeInTheDocument() + expect(screen.getByText('1/1 known results')).toBeInTheDocument() + expect(screen.queryByText('1/1 (100%)')).not.toBeInTheDocument() + }) + + it('isolates option-loading failures from the primary history request', async () => { + mockedScenariosApi.listCatalog.mockRejectedValueOnce(new Error('catalog unavailable')) + mockedScenariosApi.listRuns.mockResolvedValue({ + items: [RUN], + pagination: { limit: 25, has_more: false }, + }) + renderHistory() + + expect(await screen.findByTestId('scenario-history-table')).toBeInTheDocument() + expect(screen.getByText(/filter options could not be loaded: scenario names/i)).toBeInTheDocument() + }) + + it('shows request errors and retries without swallowing the failure', async () => { + const user = userEvent.setup() + mockedScenariosApi.listRuns + .mockRejectedValueOnce(new Error('history unavailable')) + .mockResolvedValueOnce({ + items: [RUN], + pagination: { limit: 25, has_more: false }, + }) + renderHistory() + + expect(await screen.findByTestId('scenario-history-error')).toHaveTextContent('history unavailable') + await user.click(screen.getByRole('button', { name: 'Retry' })) + expect(await screen.findByTestId('scenario-history-table')).toBeInTheDocument() + expect(mockedScenariosApi.listRuns).toHaveBeenCalledTimes(2) + }) + + it('distinguishes unfiltered and filtered empty states', async () => { + const user = userEvent.setup() + const onNavigate = jest.fn() + mockedScenariosApi.listRuns.mockResolvedValue({ + items: [], + pagination: { limit: 25, has_more: false }, + }) + const first = renderHistory({ ...defaultProps, onNavigate }) + + expect(await screen.findByText(/launch a scenario/i)).toBeInTheDocument() + await user.click(screen.getByRole('button', { name: 'Browse scenarios' })) + expect(onNavigate).toHaveBeenCalledWith('scenarios') + first.unmount() + + renderHistory({ + ...defaultProps, + filters: { ...DEFAULT_SCENARIO_HISTORY_FILTERS, statuses: ['FAILED'] }, + }) + expect(await screen.findByText('Try adjusting your filters.')).toBeInTheDocument() + expect(screen.queryByRole('button', { name: 'Browse scenarios' })).not.toBeInTheDocument() + }) + + it('serializes filters, paginates by cursor, and refreshes from the first page', async () => { + const user = userEvent.setup() + mockedScenariosApi.listRuns + .mockResolvedValueOnce({ + items: [RUN], + pagination: { limit: 25, has_more: true, next_cursor: 'next-page' }, + }) + .mockResolvedValue({ + items: [RUN], + pagination: { limit: 25, has_more: false }, + }) + const history = renderHistory({ + ...defaultProps, + filters: { + ...DEFAULT_SCENARIO_HISTORY_FILTERS, + scenarioNames: ['foundry.red_team'], + statuses: ['IN_PROGRESS', 'FAILED'], + operator: ['alice'], + operation: ['nightly'], + otherLabels: ['team:safety'], + }, + }) + + await screen.findByTestId('scenario-history-table') + expect(mockedScenariosApi.listRuns).toHaveBeenNthCalledWith(1, { + limit: 25, + cursor: undefined, + scenario_names: ['foundry.red_team'], + run_statuses: ['IN_PROGRESS', 'FAILED'], + label: ['operator:alice', 'operation:nightly', 'team:safety'], + }) + + await user.click(screen.getByRole('button', { name: 'Next' })) + await waitFor(() => expect(mockedScenariosApi.listRuns).toHaveBeenNthCalledWith( + 2, + expect.objectContaining({ cursor: 'next-page' }), + )) + expect(screen.getByText('Page 2')).toBeInTheDocument() + + history.rerender( + + + , + ) + await waitFor(() => expect(mockedScenariosApi.listRuns).toHaveBeenNthCalledWith( + 3, + expect.objectContaining({ cursor: undefined, run_statuses: ['COMPLETED'] }), + )) + expect(await screen.findByText('Page 1')).toBeInTheDocument() + + await user.click(screen.getByTestId('scenario-history-refresh')) + await waitFor(() => expect(mockedScenariosApi.listRuns).toHaveBeenNthCalledWith( + 4, + expect.objectContaining({ cursor: undefined }), + )) + }) + + it('hides stale pagination while changed filters are loading', async () => { + let resolveFilteredRequest: ((value: Awaited>) => void) | undefined + mockedScenariosApi.listRuns + .mockResolvedValueOnce({ + items: [RUN], + pagination: { limit: 25, has_more: true, next_cursor: 'stale-cursor' }, + }) + .mockImplementationOnce(() => new Promise((resolve) => { + resolveFilteredRequest = resolve + })) + + const history = renderHistory() + expect(await screen.findByRole('button', { name: 'Next' })).toBeEnabled() + + history.rerender( + + + , + ) + + expect(screen.queryByRole('button', { name: 'Next' })).not.toBeInTheDocument() + expect(screen.getByText('Loading scenario history...')).toBeInTheDocument() + await waitFor(() => expect(mockedScenariosApi.listRuns).toHaveBeenCalledTimes(2)) + expect(mockedScenariosApi.listRuns).toHaveBeenLastCalledWith( + expect.objectContaining({ cursor: undefined, run_statuses: ['FAILED'] }), + ) + + resolveFilteredRequest?.({ + items: [RUN], + pagination: { limit: 25, has_more: false }, + }) + expect(await screen.findByTestId('scenario-history-table')).toBeInTheDocument() + }) +}) diff --git a/frontend/src/components/History/ScenarioHistory.tsx b/frontend/src/components/History/ScenarioHistory.tsx new file mode 100644 index 0000000000..55c3af498e --- /dev/null +++ b/frontend/src/components/History/ScenarioHistory.tsx @@ -0,0 +1,485 @@ +import { useCallback, useEffect, useState } from 'react' + +import { + Badge, + Button, + Combobox, + MessageBar, + MessageBarBody, + mergeClasses, + Option, + Spinner, + Table, + TableBody, + TableCell, + TableHeader, + TableHeaderCell, + TableRow, + Text, + Tooltip, +} from '@fluentui/react-components' +import { + ArrowLeftRegular, + ArrowRightRegular, + ArrowSyncRegular, + FilterDismissRegular, + FilterRegular, + ScriptRegular, +} from '@fluentui/react-icons' + +import { labelsApi, scenariosApi } from '@/services/api' +import { toApiError } from '@/services/errors' +import type { ScenarioRunListItem, ScenarioRunState } from '@/types' +import { fetchAllPages } from '@/utils/fetchAllPages' + +import type { ViewName } from '../Sidebar/Navigation' +import { useScenarioHistoryStyles } from './ScenarioHistory.styles' +import { + DEFAULT_SCENARIO_HISTORY_FILTERS, + SCENARIO_RUN_STATES, + type ScenarioHistoryFilters, +} from './scenarioHistoryFilters' + +const PAGE_SIZE = 25 + +interface ScenarioHistoryProps { + filters: ScenarioHistoryFilters + onFiltersChange: (filters: ScenarioHistoryFilters) => void + onOpenRun: (scenarioResultId: string) => void + onNavigate: (view: ViewName) => void +} + +interface MultiFilterProps { + label: string + placeholder: string + selected: string[] + options: readonly string[] + onSelect: (values: string[]) => void + testId: string + className: string +} + +function MultiFilter({ + label, + placeholder, + selected, + options, + onSelect, + testId, + className, +}: MultiFilterProps) { + return ( + onSelect(data.selectedOptions)} + data-testid={testId} + > + {options.map((option) => )} + + ) +} + +export default function ScenarioHistory({ + filters, + onFiltersChange, + onOpenRun, + onNavigate, +}: ScenarioHistoryProps) { + const styles = useScenarioHistoryStyles() + const [runs, setRuns] = useState([]) + const [loading, setLoading] = useState(true) + const [error, setError] = useState(null) + const [optionsError, setOptionsError] = useState(null) + const [scenarioOptions, setScenarioOptions] = useState([]) + const [operatorOptions, setOperatorOptions] = useState([]) + const [operationOptions, setOperationOptions] = useState([]) + const [otherLabelOptions, setOtherLabelOptions] = useState([]) + const [page, setPage] = useState(0) + const [nextCursor, setNextCursor] = useState() + const [hasMore, setHasMore] = useState(false) + const filterKey = JSON.stringify([ + filters.scenarioNames, + filters.statuses, + filters.operator, + filters.operation, + filters.otherLabels, + ]) + const [settledFilterKey, setSettledFilterKey] = useState(null) + const [fetchToken, setFetchToken] = useState({ + cursor: undefined as string | undefined, + filterKey, + nonce: 0, + }) + + const requestPage = useCallback((cursor?: string) => { + setLoading(true) + setError(null) + setFetchToken((previous) => ({ cursor, filterKey, nonce: previous.nonce + 1 })) + }, [filterKey]) + + useEffect(() => { + let cancelled = false + Promise.allSettled([ + fetchAllPages((cursor) => scenariosApi.listCatalog(100, cursor)), + labelsApi.getLabels('scenarios'), + ]).then(([catalogResult, labelsResult]) => { + if (cancelled) return + const failures: string[] = [] + if (catalogResult.status === 'fulfilled') { + setScenarioOptions(catalogResult.value.map((scenario) => scenario.scenario_name).sort()) + } else { + failures.push('scenario names') + } + if (labelsResult.status === 'fulfilled') { + const operators = labelsResult.value.labels.operator ?? [] + const operations = labelsResult.value.labels.operation ?? [] + const others = Object.entries(labelsResult.value.labels) + .filter(([key]) => key !== 'operator' && key !== 'operation' && key !== 'source') + .flatMap(([key, values]) => values.map((value) => `${key}:${value}`)) + setOperatorOptions([...operators].sort()) + setOperationOptions([...operations].sort()) + setOtherLabelOptions(others.sort()) + } else { + failures.push('labels') + } + setOptionsError(failures.length > 0 ? `Some filter options could not be loaded: ${failures.join(', ')}.` : null) + }) + return () => { + cancelled = true + } + }, []) + + useEffect(() => { + let cancelled = false + const effectiveCursor = fetchToken.filterKey === filterKey ? fetchToken.cursor : undefined + const label = [ + ...filters.operator.map((value) => `operator:${value}`), + ...filters.operation.map((value) => `operation:${value}`), + ...filters.otherLabels, + ] + scenariosApi.listRuns({ + limit: PAGE_SIZE, + cursor: effectiveCursor, + scenario_names: filters.scenarioNames.length > 0 ? filters.scenarioNames : undefined, + run_statuses: filters.statuses.length > 0 ? filters.statuses : undefined, + label: label.length > 0 ? label : undefined, + }).then((response) => { + if (cancelled) return + setRuns(response.items) + setHasMore(response.pagination.has_more) + setNextCursor(response.pagination.next_cursor ?? undefined) + setSettledFilterKey(filterKey) + setError(null) + if (!effectiveCursor) setPage(0) + }).catch((requestError: unknown) => { + if (cancelled) return + setRuns([]) + setHasMore(false) + setNextCursor(undefined) + setSettledFilterKey(filterKey) + setError(toApiError(requestError).detail) + if (!effectiveCursor) setPage(0) + }).finally(() => { + if (!cancelled) setLoading(false) + }) + return () => { + cancelled = true + } + }, [ + fetchToken, + filterKey, + filters.scenarioNames, + filters.statuses, + filters.operator, + filters.operation, + filters.otherLabels, + ]) + + const setFilter = ( + key: K, + value: ScenarioHistoryFilters[K], + ): void => { + onFiltersChange({ ...filters, [key]: value }) + } + const hasFilters = filters.scenarioNames.length > 0 + || filters.statuses.length > 0 + || filters.operator.length > 0 + || filters.operation.length > 0 + || filters.otherLabels.length > 0 + const filtersPending = settledFilterKey !== filterKey + const displayLoading = loading || filtersPending + + return ( +
+
+
+ Scenario History + +
+
+ + {hasFilters && ( + + )} + setFilter('scenarioNames', values)} + testId="scenario-filter" + className={styles.filterDropdown} + /> + setFilter('statuses', values as ScenarioRunState[])} + testId="scenario-status-filter" + className={styles.filterDropdown} + /> + setFilter('operator', values)} + testId="scenario-operator-filter" + className={styles.filterDropdown} + /> + setFilter('operation', values)} + testId="scenario-operation-filter" + className={styles.filterDropdown} + /> + setFilter('otherLabels', values)} + testId="scenario-label-filter" + className={styles.filterDropdown} + /> +
+ {optionsError && ( + + {optionsError} + + )} +
+ +
+ {displayLoading ? ( +
+ ) : error ? ( +
+ {error} + +
+ ) : runs.length === 0 ? ( +
+ No scenario runs found + {hasFilters ? 'Try adjusting your filters.' : 'Launch a scenario to see its progress and results here.'} + {!hasFilters && ( + + )} +
+ ) : ( + + )} +
+ + {!displayLoading && !error && runs.length > 0 && ( +
+ + Page {page + 1} + +
+ )} +
+ ) +} + +interface ScenarioHistoryTableProps { + runs: ScenarioRunListItem[] + onOpenRun: (scenarioResultId: string) => void +} + +function ScenarioHistoryTable({ runs, onOpenRun }: ScenarioHistoryTableProps) { + const styles = useScenarioHistoryStyles() + return ( + + + + Scenario + State + Target + Created + Completed / elapsed + Work + Success + Errors / retries + Labels + + + + {runs.map((run) => ( + onOpenRun(run.scenario_result_id)} + > + + { + if (event.button !== 0 || event.metaKey || event.ctrlKey || event.shiftKey || event.altKey) { + event.stopPropagation() + return + } + event.preventDefault() + event.stopPropagation() + onOpenRun(run.scenario_result_id) + }} + > + + {run.scenario_registry_name ?? run.scenario_name} + + {run.scenario_registry_name && run.scenario_registry_name !== run.scenario_name + ? `${run.scenario_name} · v${run.scenario_version}` + : `v${run.scenario_version}`} + + + + + {formatState(run.status)} + + {run.target ? ( + +
+ {run.target.model_name ?? run.target.target_type} + + {run.target.target_type} + +
+
+ ) : 'Unavailable'} +
+ {formatTimestamp(run.created_at)} + +
+ {run.completed_at ? formatTimestamp(run.completed_at) : 'Not yet'} + {formatElapsed(run)} +
+
+ + {run.planned_total_available !== false && run.total_attacks !== null + ? `${run.completed_attacks}/${run.total_attacks}` + : `${run.completed_attacks} known / total unknown`} + + + {formatSuccess(run)} + + {run.error_attacks} / {run.total_retries} + +
+ {Object.entries(run.labels).map(([key, value]) => ( + {key}: {value} + ))} +
+
+
+ ))} +
+
+ ) +} + +function formatState(value: string): string { + return value.toLowerCase().replace(/_/g, ' ').replace(/^\w/, (letter: string) => letter.toUpperCase()) +} + +function formatTimestamp(value: string): string { + return new Date(value).toLocaleString(undefined, { + month: 'short', + day: 'numeric', + hour: '2-digit', + minute: '2-digit', + }) +} + +function formatElapsed(run: ScenarioRunListItem): string { + const start = Date.parse(run.created_at) + const end = run.completed_at ? Date.parse(run.completed_at) : Date.now() + const seconds = Math.max(0, Math.floor((end - start) / 1000)) + if (seconds < 60) return `${seconds}s elapsed` + if (seconds < 3600) return `${Math.floor(seconds / 60)}m elapsed` + return `${Math.floor(seconds / 3600)}h ${Math.floor((seconds % 3600) / 60)}m elapsed` +} + +function formatSuccess(run: ScenarioRunListItem): string { + const successful = run.successful_attacks + if (run.planned_total_available === false) { + return `${successful}/${run.completed_attacks} known results` + } + if (run.completed_attacks === 0) { + return '0/0' + } + return `${successful}/${run.completed_attacks} (${run.objective_achieved_rate}%)` +} diff --git a/frontend/src/components/History/scenarioHistoryFilters.test.ts b/frontend/src/components/History/scenarioHistoryFilters.test.ts new file mode 100644 index 0000000000..7f3ce048b1 --- /dev/null +++ b/frontend/src/components/History/scenarioHistoryFilters.test.ts @@ -0,0 +1,57 @@ +import { + DEFAULT_SCENARIO_HISTORY_FILTERS, + SCENARIO_RUN_STATES, + scenarioHistoryFiltersFromSearchParams, + scenarioHistoryFiltersToSearchParams, +} from './scenarioHistoryFilters' + +describe('scenario history URL filters', () => { + it('round-trips repeated filters and label search text', () => { + const filters = { + scenarioNames: ['red.team', 'benchmark'], + statuses: ['IN_PROGRESS', 'FAILED'] as const, + operator: ['alice', 'bob'], + operation: ['nightly'], + otherLabels: ['team:security', 'team:safety'], + labelSearchText: 'team', + } + + const params = scenarioHistoryFiltersToSearchParams({ + ...filters, + statuses: [...filters.statuses], + }) + + expect(params.getAll('scenario')).toEqual(['red.team', 'benchmark']) + expect(params.getAll('status')).toEqual(['IN_PROGRESS', 'FAILED']) + expect(scenarioHistoryFiltersFromSearchParams(params)).toEqual({ + ...filters, + statuses: [...filters.statuses], + }) + }) + + it('ignores synthetic and invalid run states without dropping valid filters', () => { + const params = new URLSearchParams('status=COMPLETED&status=QUEUED&status=UNKNOWN&operator=alice') + + expect(scenarioHistoryFiltersFromSearchParams(params)).toEqual({ + ...DEFAULT_SCENARIO_HISTORY_FILTERS, + statuses: ['COMPLETED'], + operator: ['alice'], + }) + }) + + it('round-trips every persisted run state', () => { + const filters = { + ...DEFAULT_SCENARIO_HISTORY_FILTERS, + statuses: [...SCENARIO_RUN_STATES], + } + + const params = scenarioHistoryFiltersToSearchParams(filters) + + expect(params.getAll('status')).toEqual(SCENARIO_RUN_STATES) + expect(scenarioHistoryFiltersFromSearchParams(params)).toEqual(filters) + }) + + it('omits empty filters from the URL', () => { + expect(scenarioHistoryFiltersToSearchParams(DEFAULT_SCENARIO_HISTORY_FILTERS).toString()).toBe('') + }) +}) diff --git a/frontend/src/components/History/scenarioHistoryFilters.ts b/frontend/src/components/History/scenarioHistoryFilters.ts new file mode 100644 index 0000000000..3f78640d6a --- /dev/null +++ b/frontend/src/components/History/scenarioHistoryFilters.ts @@ -0,0 +1,58 @@ +import type { ScenarioRunState } from '@/types' + +export interface ScenarioHistoryFilters { + scenarioNames: string[] + statuses: ScenarioRunState[] + operator: string[] + operation: string[] + otherLabels: string[] + labelSearchText: string +} + +export const DEFAULT_SCENARIO_HISTORY_FILTERS: ScenarioHistoryFilters = { + scenarioNames: [], + statuses: [], + operator: [], + operation: [], + otherLabels: [], + labelSearchText: '', +} + +export const SCENARIO_RUN_STATES: readonly ScenarioRunState[] = [ + 'CREATED', + 'IN_PROGRESS', + 'COMPLETED', + 'FAILED', + 'CANCELLED', +] + +const RUN_STATES = new Set(SCENARIO_RUN_STATES) + +export function scenarioHistoryFiltersFromSearchParams( + params: URLSearchParams, +): ScenarioHistoryFilters { + const statuses = params + .getAll('status') + .filter((status): status is ScenarioRunState => RUN_STATES.has(status)) + return { + scenarioNames: params.getAll('scenario'), + statuses, + operator: params.getAll('operator'), + operation: params.getAll('operation'), + otherLabels: params.getAll('label'), + labelSearchText: params.get('labelSearch') ?? '', + } +} + +export function scenarioHistoryFiltersToSearchParams( + filters: ScenarioHistoryFilters, +): URLSearchParams { + const params = new URLSearchParams() + for (const scenarioName of filters.scenarioNames) params.append('scenario', scenarioName) + for (const status of filters.statuses) params.append('status', status) + for (const operator of filters.operator) params.append('operator', operator) + for (const operation of filters.operation) params.append('operation', operation) + for (const label of filters.otherLabels) params.append('label', label) + if (filters.labelSearchText) params.set('labelSearch', filters.labelSearchText) + return params +} diff --git a/frontend/src/components/Scenarios/ScenarioRunPage.test.tsx b/frontend/src/components/Scenarios/ScenarioRunPage.test.tsx index 4c3c772308..3d34013d48 100644 --- a/frontend/src/components/Scenarios/ScenarioRunPage.test.tsx +++ b/frontend/src/components/Scenarios/ScenarioRunPage.test.tsx @@ -141,6 +141,35 @@ describe('ScenarioRunPage', () => { expect(screen.queryByRole('columnheader', { name: 'Actions' })).not.toBeInTheDocument() }) + it('renders contract-backed safe target and run configuration metadata', () => { + mockHookState(makeState({ + run: { + ...makeState().run!, + target: { + target_type: 'OpenAIChatTarget', + endpoint: 'https://example.test/v1', + model_name: 'gpt-4o', + identifier_hash: 'safe-hash', + }, + techniques_used: ['Technique One'], + datasets_used: ['harmbench'], + scenario_parameters: { max_turns: 5 }, + labels: { operator: 'alice' }, + pyrit_version: '0.10.0', + }, + })) + + renderPage() + + expect(screen.getByText('gpt-4o')).toBeInTheDocument() + expect(screen.getByText('https://example.test/v1')).toBeInTheDocument() + expect(screen.getByText('safe-hash')).toBeInTheDocument() + expect(screen.getByText('harmbench')).toBeInTheDocument() + expect(screen.getByText('max_turns: 5')).toBeInTheDocument() + expect(screen.getByText('operator: alice')).toBeInTheDocument() + expect(screen.getByText('0.10.0')).toBeInTheDocument() + }) + it('keeps legacy runs useful without misleading totals, ETA, or a progress bar', () => { mockHookState(makeState({ planComplete: false })) diff --git a/frontend/src/components/Scenarios/ScenarioRunPage.tsx b/frontend/src/components/Scenarios/ScenarioRunPage.tsx index b2516de87e..fa21c5dee4 100644 --- a/frontend/src/components/Scenarios/ScenarioRunPage.tsx +++ b/frontend/src/components/Scenarios/ScenarioRunPage.tsx @@ -33,7 +33,7 @@ import { EyeRegular, StopRegular, } from '@fluentui/react-icons' -import { Link, useNavigate, useParams } from 'react-router' +import { Link, useLocation, useNavigate, useParams } from 'react-router' import { useScenarioRunProgress } from '@/hooks/useScenarioRunProgress' import { scenariosApi } from '@/services/api' @@ -64,6 +64,7 @@ const INTERACTIVE_ELEMENT_SELECTOR = 'a, button, input, select, textarea, [role= const RUN_BADGE_COLORS: Record = { CREATED: 'informative', + QUEUED: 'informative', IN_PROGRESS: 'brand', COMPLETED: 'success', FAILED: 'danger', @@ -88,6 +89,7 @@ interface ScenarioRunPageContentProps { function ScenarioRunPageContent({ scenarioResultId }: ScenarioRunPageContentProps) { const styles = useScenarioRunPageStyles() + const location = useLocation() const navigate = useNavigate() const { state, retry, applyRunSummary } = useScenarioRunProgress(scenarioResultId) const [nowMilliseconds, setNowMilliseconds] = useState(() => Date.now()) @@ -96,6 +98,19 @@ function ScenarioRunPageContent({ scenarioResultId }: ScenarioRunPageContentProp const [cancelError, setCancelError] = useState(null) const [selectedAttempt, setSelectedAttempt] = useState(null) const detailsTriggerRef = useRef(null) + const navigationState = location.state as { + fromScenarioHistory?: boolean + scenarioHistorySearch?: string + scenarioName?: string + } | null + const backPath = navigationState?.fromScenarioHistory + ? `/scenario-history${navigationState.scenarioHistorySearch ?? ''}` + : navigationState?.scenarioName + ? `/scanner/${encodeURIComponent(navigationState.scenarioName)}` + : '/scenario-history' + const backLabel = navigationState?.scenarioName && !navigationState.fromScenarioHistory + ? 'Back to scenario' + : 'Back to scenario history' const overall = useMemo(() => getOverallProgress(state), [state]) const techniques = useMemo(() => getTechniqueRollups(state), [state]) @@ -149,8 +164,8 @@ function ScenarioRunPageContent({ scenarioResultId }: ScenarioRunPageContentProp return (
- - Back to scanners + + {backLabel}
@@ -171,8 +186,8 @@ function ScenarioRunPageContent({ scenarioResultId }: ScenarioRunPageContentProp return (
- - Back to scanners + + {backLabel}
@@ -191,8 +206,8 @@ function ScenarioRunPageContent({ scenarioResultId }: ScenarioRunPageContentProp return (
- - Back to scanners + + {backLabel}
@@ -222,8 +237,8 @@ function ScenarioRunPageContent({ scenarioResultId }: ScenarioRunPageContentProp return (
- - Back to scanners + + {backLabel}
@@ -278,8 +293,52 @@ function ScenarioRunPageContent({ scenarioResultId }: ScenarioRunPageContentProp Completed {run.completed_at ? formatTimestamp(run.completed_at) : 'Not yet'}
+ {run.target && ( +
+ Target + {run.target.model_name ?? run.target.target_type} + {run.target.target_type} +
+ )} + {run.pyrit_version && ( +
+ PyRIT version + {run.pyrit_version} +
+ )}
+
+
+ + Run configuration + + Persisted, secret-free settings for this run. +
+
+ 0 ? run.techniques_used?.join(', ') ?? '' : 'Unavailable'} + /> + 0 ? run.datasets_used?.join(', ') ?? '' : 'Unavailable'} + /> + + + {run.target?.endpoint && } + {run.target?.identifier_hash && ( + + )} +
+
+ {state.stale && ( @@ -663,6 +722,21 @@ interface MetricProps { readonly value: string } +interface ConfigurationItemProps { + readonly label: string + readonly value: string +} + +function ConfigurationItem({ label, value }: ConfigurationItemProps) { + const styles = useScenarioRunPageStyles() + return ( +
+ {label} + {value} +
+ ) +} + function Metric({ label, value }: MetricProps) { const styles = useScenarioRunPageStyles() return ( @@ -727,6 +801,7 @@ function formatTimestamp(timestamp: string): string { if (Number.isNaN(date.getTime())) { return 'Unavailable' } + return date.toLocaleString(undefined, { month: 'short', day: 'numeric', @@ -737,6 +812,16 @@ function formatTimestamp(timestamp: string): string { }) } +function formatConfiguration(value: Record): string { + const entries = Object.entries(value) + if (entries.length === 0) { + return 'None' + } + return entries + .map(([key, item]) => `${key}: ${typeof item === 'string' ? item : JSON.stringify(item)}`) + .join(', ') +} + function formatDuration(milliseconds: number): string { if (!Number.isFinite(milliseconds) || milliseconds < 0) { return 'Unavailable' diff --git a/frontend/src/components/Sidebar/Navigation.test.tsx b/frontend/src/components/Sidebar/Navigation.test.tsx index e85813e3bd..8c61db38ae 100644 --- a/frontend/src/components/Sidebar/Navigation.test.tsx +++ b/frontend/src/components/Sidebar/Navigation.test.tsx @@ -126,7 +126,7 @@ describe("Navigation", () => { ).toBeInTheDocument(); }); - it("places Scanner immediately after Attack History without a history placeholder", () => { + it("renders the final primary navigation order", () => { renderWithProvider(); const navigation = screen.getByRole("navigation", { name: "Primary" }); const labels = within(navigation) @@ -138,11 +138,28 @@ describe("Navigation", () => { "Chat", "Attack History", "Scanner", + "Scenario History", "Targets", "Initializers", "Configuration", ]); - expect(screen.queryByRole("button", { name: "Scenario History" })).not.toBeInTheDocument(); + }); + + it("marks Scenario History current and navigates to its dedicated view", async () => { + const user = userEvent.setup(); + const onNavigate = jest.fn(); + renderWithProvider( + , + ); + + const button = screen.getByRole("button", { name: "Scenario History" }); + expect(button).toHaveAttribute("aria-current", "page"); + await user.click(button); + expect(onNavigate).toHaveBeenCalledWith("scenarioHistory"); }); it("calls onNavigate with 'scenarios' when the scenarios button is clicked", async () => { diff --git a/frontend/src/components/Sidebar/Navigation.tsx b/frontend/src/components/Sidebar/Navigation.tsx index 816eb40fe7..7ba48fb620 100644 --- a/frontend/src/components/Sidebar/Navigation.tsx +++ b/frontend/src/components/Sidebar/Navigation.tsx @@ -15,6 +15,7 @@ import { HistoryRegular, PersonFeedbackRegular, ScriptRegular, + TableRegular, WrenchRegular, OpenRegular, WeatherMoonRegular, @@ -32,6 +33,7 @@ export type ViewName = | 'targets' | 'initializers' | 'configuration' + | 'scenarioHistory' | 'scenarios' interface NavigationProps { @@ -120,6 +122,16 @@ export default function Navigation({ onClick={() => onNavigate('scenarios')} /> +
)} + {queued && ( +
+ Waiting position + {run.queue_position ?? 'Updating'} +
+ )}
@@ -352,6 +364,23 @@ function ScenarioRunPageContent({ scenarioResultId }: ScenarioRunPageContentProp )} + + + {(state.overloadSummaries?.length ?? 0) > 0 && ( + + + Recent target overload detected: {formatOverloadSummaries(state.overloadSummaries ?? [])}. + {' '}PyRIT is retrying these requests without adaptive throttling; concurrency is not automatically reduced yet. + + + )} + {run.status === 'FAILED' && ( @@ -375,6 +404,28 @@ function ScenarioRunPageContent({ scenarioResultId }: ScenarioRunPageContentProp {progressText}
+ {queued ? ( +
+
+ + Queued{run.queue_position ? ` · Position ${run.queue_position}` : ''} + + + {run.active_scenario_result_id + ? `Waiting for active run ${run.active_scenario_result_id} to finish.` + : 'Waiting for the scheduler to start this run.'} + +
+
+ Execution progress + Not started +
+
+ Estimated remaining + Available after start +
+
+ ) : (
@@ -404,8 +455,9 @@ function ScenarioRunPageContent({ scenarioResultId }: ScenarioRunPageContentProp
+ )} - {isTerminalRunState(run.status) ? `Run ${formatRunState(run.status)}` : ''} + {queued ? progressText : isTerminalRunState(run.status) ? `Run ${formatRunState(run.status)}` : ''} @@ -645,7 +697,9 @@ function ScenarioRunPageContent({ scenarioResultId }: ScenarioRunPageContentProp Cancel this scenario run? - In-flight work will be stopped. Attempts already persisted will remain available in this dashboard. + {queued + ? 'This run will be removed from the queue and will never execute.' + : 'In-flight work will be stopped. Attempts already persisted will remain available in this dashboard.'} {cancelError && ( @@ -657,6 +711,7 @@ function ScenarioRunPageContent({ scenarioResultId }: ScenarioRunPageContentProp +
+ + setQuery(data.value)} + placeholder="Search datasets" + aria-label="Search datasets" + data-testid="dataset-search-input" + /> + + {status === 'loading' && ( + + Loading registered datasets… + + )} + {status === 'error' && ( + + + Registered datasets couldn’t be loaded. Scenario defaults remain available. + {error ? ` ${error}` : ''} + + + )} +
+ {visibleDatasets.length > 0 ? ( + visibleDatasets.map((name) => ( + onChange(name, data.checked === true)} + data-testid={`dataset-${name}`} + /> + )) + ) : ( + No datasets match this search. + )} +
+ + ) +} + interface ScenarioDetailProps { activeTarget: TargetInstance | null labels: Record @@ -510,8 +671,8 @@ function ScenarioDetailContent({ return (
- - Back to scanners + + Back to scenarios
Scenario "{decodedScenarioName}" was not found @@ -526,8 +687,8 @@ function ScenarioDetailContent({ return (
- - Back to scanners + + Back to scenarios
@@ -560,7 +721,6 @@ function ScenarioDetailContent({ targets={targets} activeTarget={activeTarget} labels={labels} - onNavigate={onNavigate} /> ) } @@ -570,24 +730,21 @@ interface ScenarioLaunchFormProps { targets: TargetInstance[] activeTarget: TargetInstance | null labels: Record - onNavigate: (view: ViewName) => void } -function ScenarioLaunchForm({ - scenario, - targets, - activeTarget, - labels, - onNavigate, -}: ScenarioLaunchFormProps) { +function ScenarioLaunchForm({ scenario, targets, activeTarget, labels }: ScenarioLaunchFormProps) { const styles = useScenarioDetailStyles() const navigate = useNavigate() const formId = `scenario-launch-${encodeURIComponent(scenario.scenario_name).replace(/%/g, '-')}` - const { techniques: techniqueOptions, defaultTechniques } = useMemo( + const { presets, concrete, defaultSelection, initialCustomTechniques } = useMemo( () => uniqueTechniqueOptions(scenario), [scenario], ) + const techniqueSummaries = useMemo( + () => new Map(scenario.technique_summaries.map((summary) => [summary.name, summary])), + [scenario.technique_summaries], + ) const dynamicParameters = useMemo( () => scenario.supported_parameters.filter( (parameter) => !COMMON_SCENARIO_PARAMETER_NAMES.has(parameter.name), @@ -595,88 +752,73 @@ function ScenarioLaunchForm({ [scenario.supported_parameters], ) const isBaselineForbidden = scenario.baseline_policy === 'forbidden' + const usesAdaptiveTechniqueSelection = scenario.scenario_name === TEXT_ADAPTIVE_SCENARIO_NAME const [targetName, setTargetName] = useState(() => { if (activeTarget && targets.some((target) => target.target_registry_name === activeTarget.target_registry_name)) { return activeTarget.target_registry_name } - return targets[0]?.target_registry_name ?? '' + return targets[0].target_registry_name }) - const [selectedTechniques, setSelectedTechniques] = useState(() => defaultTechniques) + const [techniqueSelection, setTechniqueSelection] = useState(() => defaultSelection) + const [customTechniques, setCustomTechniques] = useState(() => initialCustomTechniques) 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 [harmCategoriesFilter, setHarmCategoriesFilter] = useState('') - const [dataTypesFilter, setDataTypesFilter] = useState('') + const [availableDatasets, setAvailableDatasets] = useState(() => [...scenario.default_datasets]) + const [selectedDatasets, setSelectedDatasets] = useState(() => [...scenario.default_datasets]) + const [datasetCatalogStatus, setDatasetCatalogStatus] = useState('loading') + const [datasetCatalogError, setDatasetCatalogError] = useState(null) + const [maxDatasetSize, setMaxDatasetSize] = useState(() => { + const limit = scenario.dataset_size_limit + return limit.default_count !== null && limit.default_scope === limit.override_scope + ? String(limit.default_count) + : '' + }) + const [hasMaxDatasetSizeOverride, setHasMaxDatasetSizeOverride] = useState(false) + const [maxDatasetSizeInputRejected, setMaxDatasetSizeInputRejected] = useState(false) const [maxConcurrency, setMaxConcurrency] = useState(DEFAULT_MAX_CONCURRENCY) const [maxRetries, setMaxRetries] = useState(DEFAULT_MAX_RETRIES) - const [scenarioParamValues, setScenarioParamValues] = useState>(() => - getInitialFormValues(dynamicParameters), - ) + const [scenarioParamValues, setScenarioParamValues] = useState>(() => { + const initialValues = getInitialFormValues(dynamicParameters) + if (usesAdaptiveTechniqueSelection && MAX_ATTEMPTS_PARAMETER_NAME in initialValues) { + initialValues[MAX_ATTEMPTS_PARAMETER_NAME] = '' + } + return initialValues + }) const [validationError, setValidationError] = useState(null) const [apiError, setApiError] = useState(null) const [submitting, setSubmitting] = useState(false) - const [previewOpen, setPreviewOpen] = useState(false) const [estimateRequestState, setEstimateRequestState] = useState(null) - const [lastGoodEstimate, setLastGoodEstimate] = useState(null) + const [launchMaxAttemptsError, setLaunchMaxAttemptsError] = useState(null) + const [maxAttemptsInputRejected, setMaxAttemptsInputRejected] = useState(false) + const [adaptiveCandidateMetadata, setAdaptiveCandidateMetadata] = + useState(null) + const [adaptiveLimitNotice, setAdaptiveLimitNotice] = useState(null) // Synchronous guard against a double-submit racing ahead of the state update. const isSubmittingRef = useRef(false) const estimateSequenceRef = useRef(0) - - const selectableTechniques = useMemo( - () => [ - { - ...BASELINE_TECHNIQUE, - isBaseline: true, - disabled: isBaselineForbidden, - }, - ...techniqueOptions.map((technique) => ({ - ...technique, - isBaseline: false, - disabled: false, - })), - ], - [isBaselineForbidden, techniqueOptions], + const customTechniquesInitializedRef = useRef(defaultSelection.mode === 'custom') + const hasResolvedAdaptiveMetadataRef = useRef(false) + + const techniques = useMemo( + () => techniqueSelection.mode === 'preset' + ? [techniqueSelection.preset] + : customTechniques, + [customTechniques, techniqueSelection], ) - const techniques = selectedTechniques - const maxDatasetSizeOverride = maxDatasetSize.trim() - && maxDatasetSize !== configuredDefaultMaxDatasetSize - ? maxDatasetSize - : '' - const estimateResult = useMemo( - () => buildEstimateRequest({ - scenario, - targetName, - techniques, - dynamicParameters, - scenarioParamValues, - datasetOverride, - maxDatasetSize: maxDatasetSizeOverride, - harmCategoriesFilter, - dataTypesFilter, - includeBaseline: isBaselineForbidden ? false : baselineChecked, - }), - [ - baselineChecked, - datasetOverride, - dataTypesFilter, - dynamicParameters, - harmCategoriesFilter, - isBaselineForbidden, - maxDatasetSizeOverride, - scenario, - scenarioParamValues, - targetName, - techniques, - ], + const adaptiveCandidateScopeKey = useMemo( + () => JSON.stringify({ targetName, techniques }), + [targetName, techniques], ) + const knownAdaptiveCandidateMaximum = + adaptiveCandidateMetadata?.scopeKey === adaptiveCandidateScopeKey + ? adaptiveCandidateMetadata.maximum + : null + const adaptiveSelectionDisplayName = techniqueSelection.mode === 'preset' + ? techniqueSetDisplayName(scenario, techniqueSelection.preset) + : 'Custom selection' const requestResult = useMemo( () => buildRunRequest({ scenario, @@ -684,10 +826,9 @@ function ScenarioLaunchForm({ techniques, dynamicParameters, scenarioParamValues, - datasetOverride, - maxDatasetSize: maxDatasetSizeOverride, - harmCategoriesFilter, - dataTypesFilter, + selectedDatasets, + maxDatasetSize, + hasMaxDatasetSizeOverride, maxConcurrency, maxRetries, includeBaseline: isBaselineForbidden ? false : baselineChecked, @@ -695,24 +836,51 @@ function ScenarioLaunchForm({ }), [ baselineChecked, - datasetOverride, - dataTypesFilter, dynamicParameters, - harmCategoriesFilter, isBaselineForbidden, labels, maxConcurrency, - maxDatasetSizeOverride, + maxDatasetSize, + hasMaxDatasetSizeOverride, maxRetries, scenario, scenarioParamValues, + selectedDatasets, targetName, techniques, ], ) const estimateRequest = useMemo( - () => estimateResult.ok ? estimateResult.request : null, - [estimateResult], + () => { + if (!requestResult.ok || maxAttemptsInputRejected || maxDatasetSizeInputRejected) { + return null + } + const request = buildEstimateRequest(requestResult.request) + if (!usesAdaptiveTechniqueSelection || !request.scenario_params) { + return request + } + const scenarioParams = { ...request.scenario_params } + const configuredMaximum = scenarioParams[MAX_ATTEMPTS_PARAMETER_NAME] + if (knownAdaptiveCandidateMaximum === null || knownAdaptiveCandidateMaximum === 0) { + delete scenarioParams[MAX_ATTEMPTS_PARAMETER_NAME] + } else if ( + typeof configuredMaximum === 'number' + && configuredMaximum > knownAdaptiveCandidateMaximum + ) { + scenarioParams[MAX_ATTEMPTS_PARAMETER_NAME] = knownAdaptiveCandidateMaximum + } + return { + ...request, + scenario_params: Object.keys(scenarioParams).length > 0 ? scenarioParams : undefined, + } + }, + [ + knownAdaptiveCandidateMaximum, + maxAttemptsInputRejected, + maxDatasetSizeInputRejected, + requestResult, + usesAdaptiveTechniqueSelection, + ], ) const estimateRequestKey = useMemo( () => estimateRequest === null @@ -722,12 +890,36 @@ function ScenarioLaunchForm({ ) useEffect(() => { - if (estimateRequest === null || estimateRequestKey === null) { - return + let cancelled = false + datasetsApi + .listDatasets() + .then((response) => { + if (cancelled) return + setAvailableDatasets([...new Set([ + ...scenario.default_datasets, + ...response.items.map((item) => item.name), + ])]) + setDatasetCatalogStatus('success') + setDatasetCatalogError(null) + }) + .catch((err: unknown) => { + if (cancelled) return + setAvailableDatasets([...scenario.default_datasets]) + setDatasetCatalogStatus('error') + setDatasetCatalogError(toApiError(err).detail) + }) + return () => { + cancelled = true } + }, [scenario.default_datasets]) + useEffect(() => { const requestSequence = estimateSequenceRef.current + 1 estimateSequenceRef.current = requestSequence + if (estimateRequest === null || estimateRequestKey === null) { + return + } + const controller = new AbortController() const debounceTimer = window.setTimeout(() => { @@ -741,14 +933,60 @@ function ScenarioLaunchForm({ return } const result = mapScenarioRunEstimate(response, 'request') + const adaptiveDetails = + result.status === 'available' || result.status === 'conditional' + ? result.estimate.adaptiveDetails + : null + if (usesAdaptiveTechniqueSelection && adaptiveDetails) { + const maximum = adaptiveDetails.candidateTechniqueCount + const hadResolvedAdaptiveMetadata = hasResolvedAdaptiveMetadataRef.current + hasResolvedAdaptiveMetadataRef.current = true + setAdaptiveCandidateMetadata({ scopeKey: adaptiveCandidateScopeKey, maximum }) + const rawValue = scenarioParamValues[MAX_ATTEMPTS_PARAMETER_NAME] + const parsedValue = typeof rawValue === 'string' && rawValue.trim() !== '' + ? Number(rawValue) + : null + const configuredValue = parsedValue ?? adaptiveDetails.maxAttemptsPerObjective + const isDefaultReduction = parsedValue === null + if ( + maximum > 0 + && Number.isSafeInteger(configuredValue) + && configuredValue > maximum + ) { + setScenarioParamValues((current) => ({ + ...current, + [MAX_ATTEMPTS_PARAMETER_NAME]: String(maximum), + })) + setAdaptiveLimitNotice( + isDefaultReduction + ? { + scopeKey: adaptiveCandidateScopeKey, + message: `The scenario default of ${configuredValue.toLocaleString()} is reduced to ${ + maximum.toLocaleString() + } because ${adaptiveSelectionDisplayName} provides ${maximum.toLocaleString()} compatible ${ + maximum === 1 ? 'technique' : 'techniques' + } for this target.`, + validationState: 'none', + } + : hadResolvedAdaptiveMetadata + ? { + scopeKey: adaptiveCandidateScopeKey, + message: `Reduced to ${maximum.toLocaleString()} because ${ + adaptiveSelectionDisplayName + } provides ${maximum.toLocaleString()} compatible ${ + maximum === 1 ? 'technique' : 'techniques' + } for this target.`, + validationState: 'warning', + } + : null, + ) + } + } setEstimateRequestState({ status: 'resolved', requestKey: estimateRequestKey, result, }) - if (result.status === 'available' || result.status === 'conditional') { - setLastGoodEstimate(result) - } }) .catch((err: unknown) => { if ( @@ -757,10 +995,11 @@ function ScenarioLaunchForm({ ) { return } + const mappedError = mapEstimateError(err) setEstimateRequestState({ status: 'error', requestKey: estimateRequestKey, - error: toApiError(err).detail, + ...mappedError, }) }) }, ESTIMATE_DEBOUNCE_MS) @@ -769,103 +1008,244 @@ function ScenarioLaunchForm({ window.clearTimeout(debounceTimer) controller.abort() } - }, [estimateRequest, estimateRequestKey, scenario.scenario_name]) + }, [ + adaptiveCandidateScopeKey, + adaptiveSelectionDisplayName, + estimateRequest, + estimateRequestKey, + scenario.scenario_name, + scenarioParamValues, + usesAdaptiveTechniqueSelection, + ]) + + const currentResolvedEstimate = estimateRequestState?.requestKey === estimateRequestKey + && estimateRequestState.status === 'resolved' + ? estimateRequestState.result + : null + const currentResolvedRunEstimate = currentResolvedEstimate + && (currentResolvedEstimate.status === 'available' || currentResolvedEstimate.status === 'conditional') + ? currentResolvedEstimate.estimate + : null + const currentResolvedAdaptiveDetails = currentResolvedRunEstimate + ? currentResolvedRunEstimate.adaptiveDetails + : null + const adaptiveCandidateMaximum = currentResolvedAdaptiveDetails?.candidateTechniqueCount + ?? knownAdaptiveCandidateMaximum + const adaptiveCandidateAvailability = adaptiveCandidateMaximum === null + ? undefined + : `${adaptiveSelectionDisplayName} provides ${adaptiveCandidateMaximum.toLocaleString()} compatible ${ + adaptiveCandidateMaximum === 1 ? 'technique' : 'techniques' + } for this target.` + const maxAttemptsParameter = dynamicParameters.find( + (parameter) => parameter.name === MAX_ATTEMPTS_PARAMETER_NAME, + ) + const maxAttemptsDefault = Number(maxAttemptsParameter?.default ?? 3) + const adaptiveDefaultIsReduced = usesAdaptiveTechniqueSelection + && adaptiveCandidateMaximum !== null + && adaptiveCandidateMaximum > 0 + && Number.isSafeInteger(maxAttemptsDefault) + && maxAttemptsDefault > adaptiveCandidateMaximum + const adaptiveDefaultHint = adaptiveDefaultIsReduced + ? `Blank restores the bounded default of ${adaptiveCandidateMaximum.toLocaleString()} techniques per objective for this target.` + : MAX_ATTEMPTS_DEFAULT_HINT + const maxAttemptsRawValue = scenarioParamValues[MAX_ATTEMPTS_PARAMETER_NAME] + const maxAttemptsNumericValue = typeof maxAttemptsRawValue === 'string' + && maxAttemptsRawValue.trim() !== '' + ? Number(maxAttemptsRawValue) + : null + const maxAttemptsExceedsCandidateMaximum = adaptiveCandidateMaximum !== null + && maxAttemptsNumericValue !== null + && maxAttemptsNumericValue > adaptiveCandidateMaximum + const adaptiveMetadataUnavailable = usesAdaptiveTechniqueSelection + && adaptiveCandidateMaximum === null + const noAdaptiveCandidatesError = usesAdaptiveTechniqueSelection && adaptiveCandidateMaximum === 0 + ? 'No compatible techniques are available for this target. Choose a different technique set or target.' + : undefined + const maxAttemptsClientError = usesAdaptiveTechniqueSelection + ? maxAttemptsInputRejected + ? MAX_ATTEMPTS_VALIDATION_MESSAGE + : maxAttemptsValidationError(scenarioParamValues[MAX_ATTEMPTS_PARAMETER_NAME]) + : undefined + const currentEstimateError = estimateRequestState?.requestKey === estimateRequestKey + && estimateRequestState.status === 'error' + ? estimateRequestState + : null + const maxAttemptsFieldError = maxAttemptsClientError + ?? noAdaptiveCandidatesError + ?? currentEstimateError?.maxAttemptsError + ?? launchMaxAttemptsError + ?? undefined + const maxDatasetSizeFieldError = maxDatasetSizeInputRejected + ? MAX_DATASET_SIZE_VALIDATION_MESSAGE + : maxDatasetSizeValidationError(maxDatasetSize) let estimateState: ScenarioRunEstimateState - if (!estimateResult.ok) { + if (maxAttemptsFieldError || maxDatasetSizeFieldError) { + estimateState = { + status: 'unavailable', + scope: 'request', + label: CORRECT_HIGHLIGHTED_SETTING_MESSAGE, + } + } else if (!requestResult.ok) { estimateState = { status: 'unavailable', scope: 'request', label: 'Complete the required configuration to request an estimate.', - note: estimateResult.error, + note: requestResult.error, } - } else if ( - estimateRequestState?.requestKey === estimateRequestKey - && estimateRequestState.status === 'resolved' - ) { - estimateState = estimateRequestState.result - } else if ( - estimateRequestState?.requestKey === estimateRequestKey - && estimateRequestState.status === 'error' - ) { - estimateState = lastGoodEstimate - ? { - status: 'stale', - estimate: lastGoodEstimate.estimate, - label: 'Showing the last successful estimate.', - error: estimateRequestState.error, - } - : { - status: 'unavailable', - scope: 'request', - label: 'The backend estimate could not be refreshed.', - note: estimateRequestState.error, - } - } else if (lastGoodEstimate) { + } else if (currentResolvedEstimate) { + estimateState = currentResolvedEstimate + } else if (currentEstimateError) { estimateState = { - status: 'refreshing', - estimate: lastGoodEstimate.estimate, - label: 'Updating for the current configuration…', + status: 'unavailable', + scope: 'request', + label: currentEstimateError.summary, + note: currentEstimateError.note, } } else { estimateState = { status: 'loading', scope: 'request' } } + const estimateRequestBlocked = estimateRequestState?.requestKey === estimateRequestKey + && estimateRequestState.status === 'error' - const handleTechniqueChange = (technique: SelectableTechnique, checked: boolean): void => { - if (technique.isBaseline) { - setBaselineChecked(checked) + const handleTechniqueModeChange = (value: string): void => { + if (value === CUSTOM_TECHNIQUE_SET_VALUE) { + if (!customTechniquesInitializedRef.current) { + const members = techniqueSelection.mode === 'preset' + ? techniqueSetMembers(scenario, techniqueSelection.preset) + : [] + const concreteSet = new Set(concrete) + setCustomTechniques(members.filter((member) => concreteSet.has(member))) + customTechniquesInitializedRef.current = true + } + setTechniqueSelection({ mode: 'custom' }) } else { - setSelectedTechniques((current) => { - if (checked) { - return current.includes(technique.name) - ? current - : [...current, technique.name] - } - return current.filter((name) => name !== technique.name) - }) + setTechniqueSelection({ mode: 'preset', preset: value }) } setValidationError(null) } - const isTechniqueSelected = (technique: SelectableTechnique): boolean => ( - technique.isBaseline ? baselineChecked : selectedTechniques.includes(technique.name) - ) + const handleConcreteChange = (name: string, checked: boolean): void => { + setCustomTechniques((current) => { + if (checked) { + return current.includes(name) + ? current + : [...current, name] + } + return current.filter((technique) => technique !== name) + }) + setValidationError(null) + } - const handleTagChange = (tag: string): void => { - const members = selectableTechniques.filter( - (technique) => !technique.disabled && technique.tags.includes(tag), - ) - const shouldSelect = members.some((technique) => !isTechniqueSelected(technique)) - const memberNames = new Set( - members.filter((technique) => !technique.isBaseline).map((technique) => technique.name), - ) - setSelectedTechniques((current) => { - const selected = new Set(current) - for (const name of memberNames) { - if (shouldSelect) selected.add(name) - else selected.delete(name) + const handleDatasetChange = (name: string, checked: boolean): void => { + setSelectedDatasets((current) => { + if (checked) { + return current.includes(name) ? current : [...current, name] } - return techniqueOptions.map((technique) => technique.name).filter((name) => selected.has(name)) + return current.filter((dataset) => dataset !== name) }) - if (members.some((technique) => technique.isBaseline)) { - setBaselineChecked(shouldSelect) - } setValidationError(null) } const updateScenarioParam = (name: string, value: ParameterFormValue): void => { setScenarioParamValues((current) => ({ ...current, [name]: value })) + if (name === MAX_ATTEMPTS_PARAMETER_NAME) { + setMaxAttemptsInputRejected(false) + setAdaptiveLimitNotice(null) + setLaunchMaxAttemptsError(null) + setApiError(null) + } + setValidationError(null) + } + + const rejectScenarioParamInput = ( + name: string, + reason: RejectedNumberInputReason, + retainedValue: string, + ): void => { + if (name !== MAX_ATTEMPTS_PARAMETER_NAME) { + return + } + + if (reason === 'above-max' && adaptiveCandidateMaximum !== null) { + setScenarioParamValues((current) => ({ + ...current, + [name]: String(adaptiveCandidateMaximum), + })) + setMaxAttemptsInputRejected(false) + setAdaptiveLimitNotice({ + scopeKey: adaptiveCandidateScopeKey, + message: `Maximum reached: ${adaptiveCandidateAvailability}`, + validationState: 'warning', + }) + setLaunchMaxAttemptsError(null) + setApiError(null) + setValidationError(null) + return + } + setScenarioParamValues((current) => ({ ...current, [name]: retainedValue })) + setMaxAttemptsInputRejected(true) + setAdaptiveLimitNotice(null) + setLaunchMaxAttemptsError(null) + setApiError(null) + setValidationError(null) + } + + const updateMaxDatasetSize = (_name: string, value: ParameterFormValue): void => { + const nextValue = typeof value === 'string' ? value : '' + const parsed = Number(nextValue.trim()) + const limit = scenario.dataset_size_limit + const matchesRepresentableDefault = nextValue.trim() !== '' + && limit.default_count !== null + && limit.default_scope === limit.override_scope + && parsed === limit.default_count + setMaxDatasetSize(nextValue) + setHasMaxDatasetSizeOverride(nextValue.trim() !== '' && !matchesRepresentableDefault) + setMaxDatasetSizeInputRejected(false) + setValidationError(null) + setApiError(null) + } + + const rejectMaxDatasetSizeInput = ( + _name: string, + _reason: RejectedNumberInputReason, + retainedValue: string, + ): void => { + setMaxDatasetSize(retainedValue) + setMaxDatasetSizeInputRejected(true) + setValidationError(null) + setApiError(null) } - const handleLaunchConfirmed = async (): Promise => { + const restoreDefaultDatasetSize = (): void => { + const limit = scenario.dataset_size_limit + setMaxDatasetSize( + limit.default_count !== null && limit.default_scope === limit.override_scope + ? String(limit.default_count) + : '', + ) + setHasMaxDatasetSizeOverride(false) + setMaxDatasetSizeInputRejected(false) + setValidationError(null) + setApiError(null) + } + + const handleSubmit = async (): Promise => { if (isSubmittingRef.current) { return } setApiError(null) + if ( + adaptiveMetadataUnavailable + || adaptiveCandidateMaximum === 0 + || maxAttemptsExceedsCandidateMaximum + ) { + setValidationError('Wait for the compatible technique limit to update.') + return + } if (!requestResult.ok) { setValidationError(requestResult.error) - setPreviewOpen(false) return } @@ -875,12 +1255,17 @@ function ScenarioLaunchForm({ try { const summary = await scenariosApi.startRun(requestResult.request) - setPreviewOpen(false) navigate(`/scenario-history/${encodeURIComponent(summary.scenario_result_id)}`, { state: { scenarioName: scenario.scenario_name }, }) - } catch (err) { - setApiError(toApiError(err).detail) + } catch (err: unknown) { + const mappedError = mapEstimateError(err) + if (mappedError.maxAttemptsError) { + setLaunchMaxAttemptsError(mappedError.maxAttemptsError) + setApiError(null) + } else { + setApiError(mappedError.note ?? mappedError.summary) + } } finally { isSubmittingRef.current = false setSubmitting(false) @@ -889,23 +1274,51 @@ function ScenarioLaunchForm({ const handleFormSubmit = (event: FormEvent): void => { event.preventDefault() - setApiError(null) - if (!requestResult.ok) { - setValidationError(requestResult.error) - return - } - setValidationError(null) - setPreviewOpen(true) + void handleSubmit() } - const techniqueSelectionInvalid = selectedTechniques.length === 0 - 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 previewHarmCategories = parseDatasetNames(harmCategoriesFilter) - const previewDataTypes = parseDatasetNames(dataTypesFilter) + const techniqueSelectionInvalid = + techniqueSelection.mode === 'custom' && customTechniques.length === 0 + const datasetSelectionInvalid = scenario.default_datasets.length > 0 && selectedDatasets.length === 0 + const datasetsAreDefaults = sameStringSet(selectedDatasets, scenario.default_datasets) + const datasetSizePreview = formatDatasetSizePreview( + scenario.dataset_size_limit, + maxDatasetSize, + hasMaxDatasetSizeOverride, + ) + const presetMembers = techniqueSelection.mode === 'preset' + ? techniqueSetMembers(scenario, techniqueSelection.preset) + : [] + const atAdaptiveCandidateMaximum = adaptiveCandidateMaximum !== null + && adaptiveCandidateMaximum > 0 + && maxAttemptsNumericValue === adaptiveCandidateMaximum + const scopedAdaptiveLimitNotice = adaptiveLimitNotice?.scopeKey === adaptiveCandidateScopeKey + ? adaptiveLimitNotice + : null + const currentAdaptiveLimitNotice = scopedAdaptiveLimitNotice?.message + ?? (atAdaptiveCandidateMaximum && adaptiveCandidateAvailability + ? `Maximum reached: ${adaptiveCandidateAvailability}` + : undefined) + const currentAdaptiveLimitValidationState = scopedAdaptiveLimitNotice?.validationState + ?? (atAdaptiveCandidateMaximum && adaptiveCandidateAvailability ? 'warning' : 'none') + const currentBaselineCount = currentResolvedRunEstimate?.components + .filter((component) => component.isBaseline) + .reduce((sum, component) => sum + component.count, 0) + const baselineHint = isBaselineForbidden + ? 'This scenario does not support sending objectives directly without an attack technique, so a direct comparison cannot be included.' + : baselineChecked && currentBaselineCount + ? `Adds ${currentBaselineCount.toLocaleString()} direct ${ + currentBaselineCount === 1 ? 'baseline attack' : 'baseline attacks' + } for the current objectives.` + : 'Also send each selected objective directly, without an attack technique. This provides a comparison point for measuring whether the selected techniques improve results and adds one planned attack per objective.' + const adaptiveCapFeedback = currentResolvedAdaptiveDetails + ? formatAdaptiveCapFeedback({ + selectedCandidateCount: currentResolvedAdaptiveDetails.selectedCandidateTechniqueCount, + compatibleCandidateCount: currentResolvedAdaptiveDetails.candidateTechniqueCount, + limit: currentResolvedAdaptiveDetails.maxAttemptsPerObjective, + effectiveMaximum: currentResolvedAdaptiveDetails.techniquesPerObjectiveUpperBound, + }) + : undefined return (
- - Back to scanners + + Back to scenarios
{scenario.scenario_name} + + {scenario.scenario_type} · v{scenario.scenario_version} + +
@@ -943,16 +1366,6 @@ function ScenarioLaunchForm({ )} -
- -
-
Target @@ -990,354 +1403,396 @@ function ScenarioLaunchForm({ Techniques - Select individual techniques, or use a tag to select or clear all techniques with that tag. + Choose a predefined set, or choose Custom to select techniques individually. - {techniqueSelectionInvalid && ( - - Select at least one attack technique. - + {usesAdaptiveTechniqueSelection && ( + <> + + Core, Extra, Light, Multi-turn, and Single-turn reflect tags on PyRIT's registered + techniques. All is generated from the catalog; Recommended is curated for this scenario. + + + + Adaptive uses these as a candidate pool. It tracks one progress step per compatible objective. + Adaptive tries no more than the configured maximum or the compatible candidate count, whichever + is smaller, and stops after the first success. Adding techniques changes the candidate pool, not + the number of progress steps; compatibility can still change how many objectives can run. + + + )} -
- {selectableTechniques.map((technique) => { - const selected = isTechniqueSelected(technique) - return ( -
- + + handleTechniqueModeChange(data.value)} + aria-label="Technique set" + > + {presets.map((name) => ( + + ))} + handleTechniqueChange(technique, data.checked === true)} - data-testid={technique.isBaseline ? 'baseline-checkbox' : `technique-${technique.name}`} + label="Custom" + value={CUSTOM_TECHNIQUE_SET_VALUE} + disabled={submitting} + data-testid="technique-mode-custom" /> -
- {technique.description && ( - {technique.description} - )} - {technique.tags.length > 0 && ( -
- {technique.tags.map((tag) => { - const tagMembers = selectableTechniques.filter( - (candidate) => !candidate.disabled && candidate.tags.includes(tag), - ) - const tagSelected = tagMembers.length > 0 && tagMembers.every(isTechniqueSelected) - return ( - handleTagChange(tag)} - aria-label={`${tagSelected ? 'Clear' : 'Select'} ${techniqueSetName(tag)} techniques`} - > - {techniqueSetName(tag)} - - ) - })} -
- )} - {technique.disabled && ( - - This scenario does not support a baseline comparison. - - )} -
+
+
+ {techniqueSelection.mode === 'preset' && ( +
+ Included techniques + {presetMembers.length > 0 ? ( +
+ {presetMembers.map((name) => ( + {name} + ))} +
+ ) : ( + + No concrete members were supplied for this technique set. + + )}
- ) - })} + )} + {techniqueSelection.mode === 'custom' && ( + + {concrete.length > 0 ? ( +
+ {concrete.map((name) => { + const summary = techniqueSummaries.get(name) + return ( +
+ handleConcreteChange(name, data.checked === true)} + data-testid={`technique-${name}`} + /> + {summary?.description && ( + {summary.description} + )} +
+ ) + })} +
+ ) : ( + + No concrete techniques are registered for custom selection. + + )} +
+ )}
-
- - Parameters +
+ + Baseline -
- {dynamicParameters.map((parameter) => ( - - ))} - - setDatasetOverride(data.value)} - placeholder={scenario.default_datasets.join(', ') || undefined} - data-testid="dataset-override-input" - /> - - - setMaxDatasetSize(data.value)} - data-testid="max-dataset-size-input" - /> - - - setHarmCategoriesFilter(data.value)} - data-testid="harm-categories-filter-input" - /> - - - setDataTypesFilter(data.value)} - data-testid="data-types-filter-input" - /> - - - setMaxConcurrency(resolveSpinButtonValue(data, maxConcurrency))} - data-testid="max-concurrency-input" - /> - - - setMaxRetries(resolveSpinButtonValue(data, maxRetries))} - data-testid="max-retries-input" - /> - -
+ + setBaselineChecked(data.checked === true)} + data-testid="baseline-checkbox" + /> +
-
-
- - Run estimate + {dynamicParameters.length > 0 && ( +
+ + Scenario parameters - {displayedEstimateNotes && ( - -
+ )} + +
+ + Datasets + + + Choose the registered datasets that provide objectives for this run. + + { + setSelectedDatasets([...scenario.default_datasets]) + setValidationError(null) + }} + />
-
+ + + Advanced options + +
+ + {scenario.dataset_size_limit.override_scope !== 'unsupported' + && (hasMaxDatasetSizeOverride || maxDatasetSize.trim() === '') && ( + + )} + + setMaxConcurrency(resolveSpinButtonValue(data, maxConcurrency))} + data-testid="max-concurrency-input" + /> + + + setMaxRetries(resolveSpinButtonValue(data, maxRetries))} + data-testid="max-retries-input" + /> + +
+
+
+
+ + +
- - - { - if (!submitting) { - setPreviewOpen(data.open) - } - }} - > - - - Run preview - -
-
-
Target
-
{targetName}
-
-
-
Techniques
-
-
- {baselineChecked && baseline} - {selectedTechniques.map((name) => ( - {name} - ))} -
-
-
-
-
Datasets
-
-
- - {effectiveDatasets.length > 0 - ? effectiveDatasets.join(', ') - : 'No datasets declared'} - - - {previewDatasets.length > 0 ? 'Custom override' : 'Scenario defaults'} - {maxDatasetSize.trim() ? ` - capped at ${maxDatasetSize.trim()} each` : ''} - -
-
-
-
-
Dataset filters
-
- {previewHarmCategories.length > 0 || previewDataTypes.length > 0 ? ( -
- {previewHarmCategories.length > 0 && ( -
-
Harm categories
-
{previewHarmCategories.join(', ')}
-
- )} - {previewDataTypes.length > 0 && ( -
-
Data types
-
{previewDataTypes.join(', ')}
-
- )} -
- ) : ( - 'None' - )} -
-
-
-
Parameters
-
- {dynamicParameters.length > 0 ? ( -
- {dynamicParameters.map((parameter) => ( -
-
{parameter.name}
-
- {formatEffectiveParameterPreview( - parameter.name, - scenarioParamValues[parameter.name], - displayedEstimate, - )} -
-
- ))} -
- ) : ( - 'No scenario-specific parameters' - )} -
-
-
- -
- - - - -
-
-
+
+
diff --git a/frontend/src/components/Scenarios/ScenarioFlow.test.tsx b/frontend/src/components/Scenarios/ScenarioFlow.test.tsx index 1d660cf10c..03a7345597 100644 --- a/frontend/src/components/Scenarios/ScenarioFlow.test.tsx +++ b/frontend/src/components/Scenarios/ScenarioFlow.test.tsx @@ -4,7 +4,7 @@ import { FluentProvider, webLightTheme } from '@fluentui/react-components' import { MemoryRouter, Route, Routes, useLocation } from 'react-router' import { useScenarioRunProgress } from '@/hooks/useScenarioRunProgress' -import { scenariosApi, targetsApi } from '@/services/api' +import { datasetsApi, scenariosApi, targetsApi } from '@/services/api' import type { RegisteredScenario, ScenarioDefaultRunSizeEstimate, @@ -31,6 +31,9 @@ jest.mock('@/services/api', () => ({ targetsApi: { listTargets: jest.fn(), }, + datasetsApi: { + listDatasets: jest.fn(), + }, })) const mockUseScenarioRunProgress = useScenarioRunProgress as jest.Mock @@ -39,6 +42,7 @@ const mockGetScenario = scenariosApi.getScenario as jest.Mock const mockListCatalog = scenariosApi.listCatalog as jest.Mock const mockStartRun = scenariosApi.startRun as jest.Mock const mockListTargets = targetsApi.listTargets as jest.Mock +const mockListDatasets = datasetsApi.listDatasets as jest.Mock const SCENARIO_NAME = 'foundry.red_team_agent' const RUN_ID = '123e4567-e89b-12d3-a456-426614174000' @@ -57,7 +61,11 @@ const SCENARIO: RegisteredScenario = { }, all_techniques: ['crescendo'], default_datasets: ['harmbench'], - default_dataset_summaries: [], + dataset_size_limit: { + default_scope: 'none', + default_count: null, + override_scope: 'per_dataset', + }, baseline_policy: 'enabled', include_baseline_by_default: true, supported_parameters: [], @@ -65,8 +73,12 @@ const SCENARIO: RegisteredScenario = { version: 1, status: 'exact', total_attack_count: 2, + minimum_attack_count: null, + maximum_attack_count: null, + condition: null, components: [], datasets: [], + adaptive_details: null, note: null, retries_included: false, }, @@ -84,6 +96,9 @@ const ESTIMATE: ScenarioDefaultRunSizeEstimate = { version: 1, status: 'exact', total_attack_count: 2, + minimum_attack_count: null, + maximum_attack_count: null, + condition: null, components: [{ label: 'Configured attacks', count: 2, @@ -92,6 +107,7 @@ const ESTIMATE: ScenarioDefaultRunSizeEstimate = { note: null, }], datasets: [], + adaptive_details: null, note: null, retries_included: false, } @@ -162,6 +178,7 @@ describe('Scenario catalog-to-run integration', () => { items: [TARGET], pagination: { limit: 200, has_more: false }, }) + mockListDatasets.mockResolvedValue({ items: [{ name: 'harmbench' }] }) mockEstimateRun.mockResolvedValue(ESTIMATE) mockStartRun.mockResolvedValue({ scenario_result_id: RUN_ID }) mockUseScenarioRunProgress.mockReturnValue({ @@ -171,12 +188,13 @@ describe('Scenario catalog-to-run integration', () => { }) }) - it('carries one configured request from catalog detail through estimate, launch, and run hydration', async () => { + it('carries one configured request from catalog through estimate, launch, and run hydration', async () => { const user = userEvent.setup() renderFlow() - await user.click(await screen.findByRole('link', { name: SCENARIO_NAME })) + await user.click(await screen.findByRole('button', { name: 'Configure run' })) expect(await screen.findByRole('heading', { level: 1, name: SCENARIO_NAME })).toBeInTheDocument() + expect(screen.getByText('RedTeamAgentScenario · v1')).toBeInTheDocument() const expectedEstimateRequest = { target_name: TARGET.target_registry_name, @@ -189,7 +207,7 @@ describe('Scenario catalog-to-run integration', () => { expect.any(AbortSignal), )) expect(within(screen.getByRole('complementary', { name: 'Run preview' })) - .getByText('2 planned attacks')).toBeInTheDocument() + .getByRole('group', { name: '2 planned attacks.' })).toBeInTheDocument() await user.click(screen.getByTestId('launch-scenario-btn')) diff --git a/frontend/src/components/Scenarios/ScenarioRunEstimate.styles.ts b/frontend/src/components/Scenarios/ScenarioRunEstimate.styles.ts index 1edd185b6a..0fcc38296e 100644 --- a/frontend/src/components/Scenarios/ScenarioRunEstimate.styles.ts +++ b/frontend/src/components/Scenarios/ScenarioRunEstimate.styles.ts @@ -27,112 +27,74 @@ export const useScenarioRunEstimateStyles = makeStyles({ gap: tokens.spacingVerticalM, minWidth: 0, }, - detailGroup: { + calculationSection: { display: 'flex', flexDirection: 'column', - gap: tokens.spacingVerticalXS, - minWidth: 0, - }, - componentList: { - display: 'grid', gap: tokens.spacingVerticalS, - margin: 0, - padding: 0, - listStyleType: 'none', + minWidth: 0, }, - component: { + equation: { display: 'flex', - flexDirection: 'column', - gap: tokens.spacingVerticalXXS, - paddingLeft: tokens.spacingHorizontalS, - borderLeft: `${tokens.strokeWidthThick} solid ${tokens.colorNeutralStroke2}`, + alignItems: 'stretch', + flexWrap: 'wrap', + gap: `${tokens.spacingVerticalXS} ${tokens.spacingHorizontalXS}`, minWidth: 0, - overflowWrap: 'anywhere', }, - componentHeader: { - display: 'flex', + operand: { + display: 'inline-flex', alignItems: 'baseline', - justifyContent: 'space-between', - gap: tokens.spacingHorizontalS, - }, - componentCount: { - display: 'flex', - alignItems: 'center', - gap: tokens.spacingHorizontalXS, - flexShrink: 0, - fontVariantNumeric: 'tabular-nums', - }, - factorList: { - display: 'flex', flexWrap: 'wrap', - gap: `${tokens.spacingVerticalXXS} ${tokens.spacingHorizontalS}`, - margin: 0, - padding: 0, - listStyleType: 'none', - color: tokens.colorNeutralForeground2, - }, - datasetList: { - display: 'grid', - gap: tokens.spacingVerticalS, - }, - dataset: { - display: 'flex', - flexDirection: 'column', - gap: tokens.spacingVerticalXXS, + gap: tokens.spacingHorizontalXXS, + minWidth: 0, padding: `${tokens.spacingVerticalXS} ${tokens.spacingHorizontalS}`, + color: tokens.colorNeutralForeground1, backgroundColor: tokens.colorNeutralBackground3, - borderRadius: tokens.borderRadiusSmall, - minWidth: 0, + border: `${tokens.strokeWidthThin} solid ${tokens.colorNeutralStroke2}`, + borderRadius: tokens.borderRadiusMedium, overflowWrap: 'anywhere', }, - datasetHeader: { - display: 'flex', - alignItems: 'center', + resultOperand: { + display: 'inline-flex', + alignItems: 'baseline', flexWrap: 'wrap', - gap: tokens.spacingHorizontalXS, - }, - countList: { - display: 'grid', - gap: tokens.spacingVerticalXXS, - margin: 0, + gap: tokens.spacingHorizontalXXS, + minWidth: 0, + padding: `${tokens.spacingVerticalXS} ${tokens.spacingHorizontalS}`, + color: tokens.colorBrandForeground2, + backgroundColor: tokens.colorBrandBackground2, + border: `${tokens.strokeWidthThin} solid ${tokens.colorBrandStroke1}`, + borderRadius: tokens.borderRadiusMedium, + overflowWrap: 'anywhere', }, - countRow: { - display: 'grid', - gridTemplateColumns: 'minmax(0, 1fr) auto', - gap: tokens.spacingHorizontalS, + operandValue: { fontVariantNumeric: 'tabular-nums', - '& dd': { - margin: 0, - fontWeight: tokens.fontWeightSemibold, - }, }, - capGroup: { - display: 'flex', - flexDirection: 'column', - gap: tokens.spacingVerticalXXS, + operandDetail: { + flexBasis: '100%', + color: tokens.colorNeutralForeground3, }, - capList: { - display: 'grid', - gap: tokens.spacingVerticalXXS, - margin: 0, - paddingLeft: tokens.spacingHorizontalL, + operator: { + display: 'inline-flex', + alignItems: 'center', + minHeight: '2rem', + color: tokens.colorNeutralForeground2, + fontSize: tokens.fontSizeBase400, }, - formula: { - display: 'block', - padding: `${tokens.spacingVerticalXS} ${tokens.spacingHorizontalS}`, - overflowWrap: 'anywhere', - fontFamily: tokens.fontFamilyMonospace, - fontSize: tokens.fontSizeBase200, - backgroundColor: tokens.colorNeutralBackground3, - borderRadius: tokens.borderRadiusSmall, + calculationContext: { + maxWidth: '72ch', + color: tokens.colorNeutralForeground2, + }, + sources: { + display: 'grid', + gap: tokens.spacingVerticalXS, + paddingTop: tokens.spacingVerticalXS, + borderTop: `${tokens.strokeWidthThin} solid ${tokens.colorNeutralStroke2}`, }, - staleNotice: { + source: { display: 'flex', flexDirection: 'column', gap: tokens.spacingVerticalXXS, - padding: `${tokens.spacingVerticalXS} ${tokens.spacingHorizontalS}`, - color: tokens.colorPaletteDarkOrangeForeground1, - backgroundColor: tokens.colorPaletteDarkOrangeBackground1, - borderRadius: tokens.borderRadiusSmall, + minWidth: 0, + overflowWrap: 'anywhere', }, }) diff --git a/frontend/src/components/Scenarios/ScenarioRunEstimate.test.tsx b/frontend/src/components/Scenarios/ScenarioRunEstimate.test.tsx index 765f89d453..4ecf198111 100644 --- a/frontend/src/components/Scenarios/ScenarioRunEstimate.test.tsx +++ b/frontend/src/components/Scenarios/ScenarioRunEstimate.test.tsx @@ -1,9 +1,9 @@ import type { ReactNode } from 'react' -import { render, screen } from '@testing-library/react' +import { render, screen, within } from '@testing-library/react' import { FluentProvider, webLightTheme } from '@fluentui/react-components' -import type { ScenarioRunEstimateState, ScenarioRunSizeEstimateResponse } from '@/types' +import type { ScenarioDefaultRunSizeEstimate, ScenarioRunEstimateState } from '@/types' import { ScenarioRunEstimateDetails, @@ -15,126 +15,714 @@ function TestWrapper({ children }: { children: ReactNode }) { return {children} } -const EXACT_ESTIMATE: ScenarioRunSizeEstimateResponse = { - estimated_attack_count: 8, - components: [ - { - label: 'Prompt sending', - count: 6, - is_baseline: false, - note: 'One planned attack per selected objective and template.', - }, - { - label: 'Baseline attack', - count: 2, - is_baseline: true, - note: null, +function makeEstimate( + overrides: Partial = {}, +): ScenarioDefaultRunSizeEstimate { + return { + version: 1, + status: 'exact', + total_attack_count: 16, + minimum_attack_count: null, + maximum_attack_count: null, + condition: null, + components: [ + { + label: 'Default technique sweep', + count: 16, + factors: [ + { label: 'selected logical seed groups', count: 4 }, + { label: 'default concrete techniques', count: 4 }, + ], + is_baseline: false, + note: null, + }, + ], + datasets: [ + { + name: 'harmbench', + kind: 'dataset', + logical_seed_group_count: 400, + selected_seed_group_count: 4, + configured_caps: [], + selection_note: 'The default selection uses 4 of 400 available objectives.', + }, + ], + adaptive_details: null, + note: null, + retries_included: false, + ...overrides, + } +} + +function renderDetails(estimate: ScenarioDefaultRunSizeEstimate): void { + render( + + + , + ) +} + +describe('ScenarioRunEstimate', () => { + it('renders an exact technique-by-objective equation with a complete accessible sentence', () => { + renderDetails(makeEstimate()) + + const equation = screen.getByRole('group', { + name: '4 techniques multiplied by 4 objectives equals 16 planned attacks.', + }) + expect(within(equation).getAllByText('4')).toHaveLength(2) + expect(within(equation).getByText('techniques')).toBeInTheDocument() + expect(within(equation).getByText('objectives')).toBeInTheDocument() + expect(within(equation).getByText('16')).toBeInTheDocument() + expect(within(equation).getByText('planned attacks')).toBeInTheDocument() + expect(screen.getByText('4 objectives from harmbench · 400 available')).toBeInTheDocument() + expect(screen.getByRole('heading', { name: 'Run calculation' })).toBeInTheDocument() + }) + + it.each([1, 4, 5])( + 'renders a universal per-dataset cap of %i once before the objective-source rows', + (capCount) => { + renderDetails(makeEstimate({ + datasets: [ + { + name: 'airt_hate', + kind: 'dataset', + logical_seed_group_count: 4, + selected_seed_group_count: 4, + configured_caps: [{ + label: 'per-dataset cap', + count: capCount, + configured_on: 'dataset', + dataset_name: 'airt_hate', + }], + selection_note: null, + }, + { + name: 'airt_leakage', + kind: 'dataset', + logical_seed_group_count: 9, + selected_seed_group_count: 5, + configured_caps: [{ + label: 'per-dataset cap', + count: capCount, + configured_on: 'dataset', + dataset_name: 'airt_leakage', + }], + selection_note: null, + }, + ], + })) + + const capText = `Per-dataset cap: ${capCount} ${capCount === 1 ? 'objective' : 'objectives'}` + expect(screen.getAllByText(capText)).toHaveLength(1) + expect(screen.getByText('4 objectives from airt_hate')).toBeInTheDocument() + expect(screen.getByText('5 objectives from airt_leakage · 9 available')).toBeInTheDocument() + const sources = screen.getByRole('group', { name: 'Objective sources' }) + const sourceText = sources.textContent ?? '' + expect(sourceText.indexOf(capText)).toBeLessThan(sourceText.indexOf('4 objectives from airt_hate')) }, - ], - datasets: [ - { - name: 'harmbench', - kind: 'dataset', - logical_seed_group_count: 4, - selected_seed_group_count: 4, - configured_caps: [ - { - label: 'Jailbreak templates', + ) + + it('keeps differing dataset caps with their affected objective-source rows', () => { + renderDetails(makeEstimate({ + datasets: [ + { + name: 'dataset_alpha', + kind: 'dataset', + logical_seed_group_count: 8, + selected_seed_group_count: 4, + configured_caps: [{ + label: 'per-dataset cap', + count: 4, + configured_on: 'dataset', + dataset_name: 'dataset_alpha', + }], + selection_note: null, + }, + { + name: 'dataset_beta', + kind: 'dataset', + logical_seed_group_count: 10, + selected_seed_group_count: 5, + configured_caps: [{ + label: 'per-dataset cap', + count: 5, + configured_on: 'dataset', + dataset_name: 'dataset_beta', + }], + selection_note: null, + }, + ], + })) + + const alpha = screen.getByRole('group', { name: 'Objective source: dataset_alpha' }) + const beta = screen.getByRole('group', { name: 'Objective source: dataset_beta' }) + expect(within(alpha).getByText('Per-dataset cap: 4 objectives')).toBeInTheDocument() + expect(within(beta).getByText('Per-dataset cap: 5 objectives')).toBeInTheDocument() + expect(screen.getAllByText(/Per-dataset cap:/)).toHaveLength(2) + }) + + it('keeps a single-dataset cap attached to its objective-source row', () => { + renderDetails(makeEstimate({ + datasets: [{ + name: 'harmbench', + kind: 'dataset', + logical_seed_group_count: 8, + selected_seed_group_count: 4, + configured_caps: [{ + label: 'per-dataset cap', + count: 4, + configured_on: 'dataset', + dataset_name: 'harmbench', + }], + selection_note: null, + }], + })) + + const source = screen.getByRole('group', { name: 'Objective source: harmbench' }) + expect(within(source).getByText('Per-dataset cap: 4 objectives')).toBeInTheDocument() + }) + + it('renders a global cap once while preserving differing per-dataset caps on rows', () => { + renderDetails(makeEstimate({ + datasets: [ + { + name: 'dataset_alpha', + kind: 'dataset', + logical_seed_group_count: 8, + selected_seed_group_count: 3, + configured_caps: [ + { + label: 'per-dataset cap', + count: 3, + configured_on: 'dataset', + dataset_name: 'dataset_alpha', + }, + { + label: 'combined compound cap', + count: 10, + configured_on: 'compound', + dataset_name: null, + }, + ], + selection_note: null, + }, + { + name: 'dataset_beta', + kind: 'dataset', + logical_seed_group_count: 8, + selected_seed_group_count: 4, + configured_caps: [ + { + label: 'per-dataset cap', + count: 4, + configured_on: 'dataset', + dataset_name: 'dataset_beta', + }, + { + label: 'combined compound cap', + count: 10, + configured_on: 'compound', + dataset_name: null, + }, + ], + selection_note: null, + }, + ], + })) + + expect(screen.getAllByText('Combined compound cap: 10')).toHaveLength(1) + expect(within(screen.getByRole('group', { + name: 'Objective source: dataset_alpha', + })).getByText('Per-dataset cap: 3 objectives')).toBeInTheDocument() + expect(within(screen.getByRole('group', { + name: 'Objective source: dataset_beta', + })).getByText('Per-dataset cap: 4 objectives')).toBeInTheDocument() + }) + + it('keeps a configuration cap on only the rows where it applies', () => { + renderDetails(makeEstimate({ + datasets: [ + { + name: 'dataset_alpha', + kind: 'dataset', + logical_seed_group_count: 8, + selected_seed_group_count: 4, + configured_caps: [{ + label: 'shared configuration cap', + count: 6, + configured_on: 'configuration', + dataset_name: null, + }], + selection_note: null, + }, + { + name: 'dataset_beta', + kind: 'dataset', + logical_seed_group_count: 8, + selected_seed_group_count: 4, + configured_caps: [{ + label: 'shared configuration cap', + count: 6, + configured_on: 'configuration', + dataset_name: null, + }], + selection_note: null, + }, + { + name: 'dataset_gamma', + kind: 'dataset', + logical_seed_group_count: 8, + selected_seed_group_count: 8, + configured_caps: [], + selection_note: null, + }, + ], + })) + + expect(within(screen.getByRole('group', { + name: 'Objective source: dataset_alpha', + })).getByText('Shared configuration cap: 6')).toBeInTheDocument() + expect(within(screen.getByRole('group', { + name: 'Objective source: dataset_beta', + })).getByText('Shared configuration cap: 6')).toBeInTheDocument() + expect(within(screen.getByRole('group', { + name: 'Objective source: dataset_gamma', + })).queryByText(/Shared configuration cap/)).not.toBeInTheDocument() + expect(screen.getAllByText('Shared configuration cap: 6')).toHaveLength(2) + }) + + it('renders no cap summary for uncapped datasets and preserves the selection note', () => { + renderDetails(makeEstimate({ + datasets: [{ + name: 'harmbench', + kind: 'dataset', + logical_seed_group_count: 8, + selected_seed_group_count: 4, + configured_caps: [], + selection_note: 'Four compatible objectives remain after filtering.', + }], + })) + + expect(screen.queryByText(/cap:/i)).not.toBeInTheDocument() + expect(screen.getByText('4 objectives from harmbench · 8 available')).toBeInTheDocument() + expect(screen.getByText('Four compatible objectives remain after filtering.')).toBeInTheDocument() + }) + + it('renders heterogeneous compatibility as truthful per-technique additive terms', () => { + renderDetails(makeEstimate({ + total_attack_count: 6, + components: [ + { + label: 'technique_alpha', + count: 4, + factors: [ + { label: 'selected concrete techniques', count: 1 }, + { label: 'compatible logical seed groups', count: 4 }, + ], + is_baseline: false, + note: null, + }, + { + label: 'technique_beta', count: 2, - configured_on: 'configuration', - dataset_name: null, + factors: [ + { label: 'selected concrete techniques', count: 1 }, + { label: 'compatible logical seed groups', count: 2 }, + ], + is_baseline: false, + note: null, }, ], - selection_note: 'Four compatible objective groups selected.', - }, - ], - note: 'The backend total is authoritative.', -} + })) -describe('ScenarioRunEstimate', () => { - it('renders only the authoritative total and formula in the detailed preview', () => { - const state = mapScenarioRunEstimate(EXACT_ESTIMATE, 'request') + const equation = screen.getByTestId('run-calculation') + expect(within(equation).getByText('objectives · Technique alpha')).toBeInTheDocument() + expect(within(equation).getByText('objectives · Technique beta')).toBeInTheDocument() + expect(equation).toHaveTextContent('4objectives · Technique alpha+2objectives · Technique beta=6planned attacks') + }) + + it('uses parentheses to make baseline precedence explicit', () => { + renderDetails(makeEstimate({ + total_attack_count: 20, + components: [ + ...makeEstimate().components, + { + label: 'Baseline', + count: 4, + factors: [{ label: 'selected logical seed groups', count: 4 }], + is_baseline: true, + note: null, + }, + ], + })) + + const equation = screen.getByTestId('run-calculation') + expect(equation).toHaveTextContent('(4techniques×4objectives)+4direct baseline attacks=20planned attacks') + expect(screen.getByRole('group', { + name: '( 4 techniques multiplied by 4 objectives ) plus 4 direct baseline attacks equals 20 planned attacks.', + })).toBeInTheDocument() + }) + + it('keeps guaranteed and target-conditional terms in one bounded equation', () => { + renderDetails(makeEstimate({ + status: 'conditional', + total_attack_count: null, + minimum_attack_count: 12, + maximum_attack_count: 20, + condition: 'target_capabilities', + components: [ + { + label: 'Baseline', + count: 4, + factors: [{ label: 'objectives', count: 4 }], + is_baseline: true, + note: null, + }, + { + label: 'Inline jailbreak delivery', + count: 8, + factors: [ + { label: 'objectives', count: 4 }, + { label: 'jailbreak templates', count: 2 }, + ], + is_baseline: false, + note: null, + }, + { + label: 'Native system-prompt jailbreak delivery', + count: 8, + factors: [ + { label: 'objectives', count: 4 }, + { label: 'jailbreak templates', count: 2 }, + ], + is_baseline: false, + condition: 'target_capabilities', + note: null, + }, + ], + })) + + const equation = screen.getByTestId('run-calculation') + expect(equation).toHaveTextContent('4objectives · Inline jailbreak delivery') + expect(equation).toHaveTextContent('4objectives · Native system-prompt jailbreak delivery · if supported') + expect(equation).toHaveTextContent('4direct baseline attacks') + expect(equation).toHaveTextContent('12–20planned attacks') + }) + it('shows adaptive progress objectives and the bounded underlying attempt work', () => { + const estimate = makeEstimate({ + status: 'conditional', + total_attack_count: null, + components: [], + adaptive_details: { + objective_count: 21, + selected_candidate_technique_count: 2, + candidate_technique_count: 2, + max_attempts_per_objective: 3, + techniques_per_objective_upper_bound: 2, + technique_attempt_count_upper_bound: 42, + stop_on_first_success: true, + compatibility_may_reduce_attempts: true, + }, + }) + const state = mapScenarioRunEstimate(estimate, 'request') render( + , ) - expect(screen.getByText('8 attacks')).toBeInTheDocument() + expect(screen.getByText('21 objectives · up to 42 technique attempts')).toBeInTheDocument() + expect(screen.getByRole('group', { + name: '21 objectives multiplied by up to 2 techniques per objective, the smaller of 2 selected candidates and limit 3, equals up to 42 technique attempts.', + })).toBeInTheDocument() + expect(screen.getByText('2 selected candidates · limit 3')).toBeInTheDocument() + expect(screen.getByRole('group', { + name: 'Direct baseline comparison is not included: up to 21 Adaptive attacks. Planned total is confirmed at launch.', + })).toBeInTheDocument() + expect(screen.queryByText('Exact total')).not.toBeInTheDocument() expect(screen.getByText( - 'Prompt sending: 6 + Baseline attack: 2 = 8', + 'Technique-attempt totals exclude multi-turn target exchanges and retries. Adaptive stops each objective after the first successful technique. Compatibility may reduce how many candidates each objective can try.', )).toBeInTheDocument() - expect(screen.queryByText('Backend estimate')).not.toBeInTheDocument() - expect(screen.queryByText('Current configuration')).not.toBeInTheDocument() - expect(screen.queryByText('Planned components')).not.toBeInTheDocument() - expect(screen.queryByText('Dataset populations')).not.toBeInTheDocument() }) - it('supports loading, conditional null totals, unavailable, and stale states', () => { - const loading: ScenarioRunEstimateState = { status: 'loading', scope: 'request' } - const { rerender } = render( + it('shows baseline-aware planned attacks before unchanged Adaptive work', () => { + const estimate = makeEstimate({ + status: 'conditional', + total_attack_count: null, + minimum_attack_count: 21, + maximum_attack_count: 42, + components: [ + { + label: 'Baseline', + count: 21, + factors: [{ label: 'objectives', count: 21 }], + is_baseline: true, + note: null, + }, + { + label: 'Adaptive objectives', + count: 21, + factors: [{ label: 'compatible objectives', count: 21 }], + is_baseline: false, + note: null, + }, + ], + adaptive_details: { + objective_count: 21, + selected_candidate_technique_count: 2, + candidate_technique_count: 2, + max_attempts_per_objective: 3, + techniques_per_objective_upper_bound: 2, + technique_attempt_count_upper_bound: 42, + stop_on_first_success: true, + compatibility_may_reduce_attempts: true, + }, + }) + const state = mapScenarioRunEstimate(estimate, 'request') + render( - + + , ) - expect(screen.getByText('Calculating run estimate...')).toBeInTheDocument() - const conditional = mapScenarioRunEstimate({ - ...EXACT_ESTIMATE, - estimated_attack_count: null, - minimum_attack_count: 12, - maximum_attack_count: 20, - components: [{ label: 'Possible attacks', count: 20, is_baseline: false, note: null }], - datasets: [], - note: null, - }, 'default') - rerender( - - - , + expect(screen.getByText('21–42 planned attacks · up to 42 technique attempts')).toBeInTheDocument() + const plannedEquation = screen.getByRole('group', { + name: 'Direct baseline comparison is included: 21 direct baseline attacks plus up to 21 Adaptive attacks equals 21–42 planned attacks.', + }) + expect(plannedEquation).toHaveTextContent( + '21direct baseline attacks+up to 21Adaptive attacks=21–42planned attacks', ) - expect(screen.getByText('12-20 attacks')).toBeInTheDocument() - expect(screen.getByText( - 'Possible attacks: 20 = conditional total', - )).toBeInTheDocument() + expect(screen.getByRole('heading', { name: 'Planned attacks' })).toBeInTheDocument() + expect(screen.getByRole('heading', { name: 'Adaptive work' })).toBeInTheDocument() + const adaptiveWork = screen.getByTestId('adaptive-work-calculation') + expect(adaptiveWork).toHaveTextContent('21objectives×up to 2techniques per objective') + expect(adaptiveWork).toHaveTextContent('=up to 42technique attempts') + expect(screen.queryByText(/Attempt ceiling:/)).not.toBeInTheDocument() + expect(screen.queryByText(/Progress tracks/)).not.toBeInTheDocument() + expect(screen.queryByText(/objective envelope|logical seed groups|selected seed groups/i)).not.toBeInTheDocument() + }) - const unavailable = mapScenarioRunEstimate({ - ...EXACT_ESTIMATE, - estimated_attack_count: null, + it('removes the baseline term while keeping Adaptive work unchanged', () => { + renderDetails(makeEstimate({ + status: 'conditional', + total_attack_count: null, + minimum_attack_count: null, + maximum_attack_count: 21, + components: [ + { + label: 'Adaptive objectives', + count: 21, + factors: [{ label: 'compatible objectives', count: 21 }], + is_baseline: false, + note: null, + }, + ], + adaptive_details: { + objective_count: 21, + selected_candidate_technique_count: 14, + candidate_technique_count: 14, + max_attempts_per_objective: 14, + techniques_per_objective_upper_bound: 14, + technique_attempt_count_upper_bound: 294, + stop_on_first_success: true, + compatibility_may_reduce_attempts: true, + }, + })) + + const plannedEquation = screen.getByRole('group', { + name: 'Direct baseline comparison is not included: up to 21 Adaptive attacks equals up to 21 planned attacks.', + }) + expect(plannedEquation).toHaveTextContent('up to 21Adaptive attacks=up to 21planned attacks') + expect(within(plannedEquation).queryByText(/baseline attack/)).not.toBeInTheDocument() + const adaptiveWork = screen.getByTestId('adaptive-work-calculation') + expect(adaptiveWork).toHaveTextContent('21objectives×up to 14techniques per objective') + expect(adaptiveWork).toHaveTextContent('=up to 294technique attempts') + }) + + it('renders exact Adaptive planned values without inventing a range', () => { + renderDetails(makeEstimate({ + status: 'exact', + total_attack_count: 42, minimum_attack_count: null, maximum_attack_count: null, + components: [ + { + label: 'Baseline', + count: 21, + factors: [{ label: 'objectives', count: 21 }], + is_baseline: true, + note: null, + }, + { + label: 'Adaptive objectives', + count: 21, + factors: [{ label: 'objectives', count: 21 }], + is_baseline: false, + note: null, + }, + ], + adaptive_details: { + objective_count: 21, + selected_candidate_technique_count: 14, + candidate_technique_count: 14, + max_attempts_per_objective: 14, + techniques_per_objective_upper_bound: 14, + technique_attempt_count_upper_bound: 294, + stop_on_first_success: true, + compatibility_may_reduce_attempts: false, + }, + })) + + expect(screen.getByRole('group', { + name: 'Direct baseline comparison is included: 21 direct baseline attacks plus 21 Adaptive attacks equals 42 planned attacks.', + })).toBeInTheDocument() + expect(screen.queryByText('21–42')).not.toBeInTheDocument() + }) + + it('preserves a nonzero Adaptive planned range when no baseline is included', () => { + renderDetails(makeEstimate({ + status: 'conditional', + total_attack_count: null, + minimum_attack_count: 5, + maximum_attack_count: 21, + components: [], + adaptive_details: { + objective_count: 21, + selected_candidate_technique_count: 2, + candidate_technique_count: 2, + max_attempts_per_objective: 2, + techniques_per_objective_upper_bound: 2, + technique_attempt_count_upper_bound: 42, + stop_on_first_success: true, + compatibility_may_reduce_attempts: true, + }, + })) + + expect(screen.getByRole('group', { + name: 'Direct baseline comparison is not included: 5–21 Adaptive attacks equals 5–21 planned attacks.', + })).toBeInTheDocument() + }) + + it('uses the configured max when it is lower than the adaptive candidate pool', () => { + renderDetails(makeEstimate({ + status: 'conditional', + total_attack_count: null, + components: [], + adaptive_details: { + objective_count: 21, + selected_candidate_technique_count: 14, + candidate_technique_count: 5, + max_attempts_per_objective: 3, + techniques_per_objective_upper_bound: 3, + technique_attempt_count_upper_bound: 63, + stop_on_first_success: true, + compatibility_may_reduce_attempts: true, + }, + })) + + expect(screen.getByText('5 compatible candidates from 14 selected · limit 3')).toBeInTheDocument() + expect(screen.getByText('up to 63')).toBeInTheDocument() + }) + + it('uses the candidate pool when it is lower than the adaptive max', () => { + renderDetails(makeEstimate({ + status: 'conditional', + total_attack_count: null, + components: [], + adaptive_details: { + objective_count: 21, + selected_candidate_technique_count: 2, + candidate_technique_count: 2, + max_attempts_per_objective: 5, + techniques_per_objective_upper_bound: 2, + technique_attempt_count_upper_bound: 42, + stop_on_first_success: true, + compatibility_may_reduce_attempts: true, + }, + })) + + expect(screen.getByText('techniques per objective')).toBeInTheDocument() + expect(screen.getByText('2 selected candidates · limit 5')).toBeInTheDocument() + expect(screen.getByText('up to 42')).toBeInTheDocument() + }) + + it('adapts legacy version-one payloads without the selected candidate count', () => { + const estimate = mapScenarioRunEstimate(makeEstimate({ + status: 'conditional', + total_attack_count: null, + components: [], + adaptive_details: { + objective_count: 21, + candidate_technique_count: 2, + max_attempts_per_objective: 3, + techniques_per_objective_upper_bound: 2, + technique_attempt_count_upper_bound: 42, + stop_on_first_success: true, + compatibility_may_reduce_attempts: true, + }, + }), 'request') + + expect(estimate).toMatchObject({ + status: 'conditional', + estimate: { + adaptiveDetails: { + selectedCandidateTechniqueCount: 2, + }, + }, + }) + }) + + it('preserves loading, unavailable, and unknown conditional states', () => { + const loading: ScenarioRunEstimateState = { status: 'loading', scope: 'request' } + const { rerender } = render( + , + ) + expect(screen.getByText('Calculating planned attacks...')).toBeInTheDocument() + + const unavailable = mapScenarioRunEstimate(makeEstimate({ + status: 'unavailable', + total_attack_count: null, components: [], datasets: [], note: 'Target capability is not available.', - }, 'request') + }), 'request') + rerender() + expect(screen.getByText('Estimate unavailable')).toBeInTheDocument() + expect(screen.getByText('Configured run size unavailable')).toBeInTheDocument() + rerender( - - + , ) - expect(screen.getByText('Estimate unavailable')).toBeInTheDocument() - expect(screen.getByText('Configured run size unavailable')).toBeInTheDocument() - expect(screen.getByText('Target capability is not available.')).toBeInTheDocument() - - const exact = mapScenarioRunEstimate(EXACT_ESTIMATE, 'request') - if (exact.status !== 'available') { - throw new Error('Expected exact estimate to map to an available state.') - } - const stale: ScenarioRunEstimateState = { - status: 'stale', - estimate: exact.estimate, - label: 'Showing the last successful estimate.', - error: 'Preview service timed out.', - } - rerender( + expect(screen.getByText('Exact total')).toBeInTheDocument() + expect(screen.getByText('unavailable')).toBeInTheDocument() + }) + + it('does not render implementation terminology in the shared estimate surfaces', () => { + const state = mapScenarioRunEstimate(makeEstimate(), 'request') + render( - + + , ) - expect(screen.getByText('8 attacks')).toBeInTheDocument() - expect(screen.queryByText('Showing the last successful estimate.')).not.toBeInTheDocument() - expect(screen.queryByText('Preview service timed out.')).not.toBeInTheDocument() + + expect(screen.queryByText(/logical seed groups/i)).not.toBeInTheDocument() + expect(screen.queryByText(/selected seed groups/i)).not.toBeInTheDocument() + expect(screen.queryByText(/planned components/i)).not.toBeInTheDocument() + expect(screen.queryByText(/objective envelopes/i)).not.toBeInTheDocument() + expect(screen.queryByText(/how this count is calculated/i)).not.toBeInTheDocument() }) }) diff --git a/frontend/src/components/Scenarios/ScenarioRunEstimate.tsx b/frontend/src/components/Scenarios/ScenarioRunEstimate.tsx index 70f269acb0..f15d4d7583 100644 --- a/frontend/src/components/Scenarios/ScenarioRunEstimate.tsx +++ b/frontend/src/components/Scenarios/ScenarioRunEstimate.tsx @@ -3,18 +3,49 @@ import { Badge, Spinner, Text } from '@fluentui/react-components' import type { ScenarioRunEstimate, ScenarioRunEstimateComponent, + ScenarioRunEstimateDatasetCap, + ScenarioRunEstimateFactor, ScenarioRunEstimateState, } from '@/types' +import { + formatAdaptiveCapAccessibleRule, + formatAdaptiveCapMetadata, +} from './scenarioAdaptiveCap' +import { normalizeDatasetCaps } from './scenarioDatasetCaps' import { useScenarioRunEstimateStyles } from './ScenarioRunEstimate.styles' interface ScenarioRunEstimateSummaryProps { state: ScenarioRunEstimateState - compact?: boolean } interface ScenarioRunEstimateDetailsProps { state: ScenarioRunEstimateState + idPrefix?: string +} + +interface CalculationOperand { + id: string + value: string + label: string + detail?: string + result?: boolean +} + +interface CalculationOperator { + id: string + symbol: '(' | ')' | '×' | '+' | '=' +} + +type CalculationPart = + | { kind: 'operand'; operand: CalculationOperand } + | { kind: 'operator'; operator: CalculationOperator } + +interface RunCalculation { + parts: CalculationPart[] + accessibleLabel: string + summary?: string + context?: string } function stateEstimate(state: ScenarioRunEstimateState): ScenarioRunEstimate | undefined { @@ -36,18 +67,17 @@ function scopeLabel(state: ScenarioRunEstimateState): string { return scope === 'default' ? 'Default configuration' : 'Current configuration' } -function statusLabel(state: ScenarioRunEstimateState): string { +function statusLabel(state: ScenarioRunEstimateState): string | null { switch (state.status) { case 'loading': return 'Loading estimate' case 'available': - return 'Backend estimate' case 'conditional': - return 'Conditional estimate' + return null case 'refreshing': - return 'Updating estimate' + return state.label case 'stale': - return 'Previous estimate' + return 'Estimate may be out of date' case 'unavailable': return 'Estimate unavailable' } @@ -56,7 +86,6 @@ function statusLabel(state: ScenarioRunEstimateState): string { function statusColor(state: ScenarioRunEstimateState): 'brand' | 'warning' | 'subtle' { switch (state.status) { case 'available': - case 'refreshing': return 'brand' case 'conditional': case 'stale': @@ -66,6 +95,10 @@ function statusColor(state: ScenarioRunEstimateState): 'brand' | 'warning' | 'su } } +function formatCount(value: number): string { + return value.toLocaleString() +} + function formatEstimateValue(value: number): string { return value.toLocaleString() } @@ -76,63 +109,501 @@ function countLabel(value: number, singular: string, plural: string): string { function formatPlannedAttackSummary(estimate: ScenarioRunEstimate): string { if (estimate.total !== null) { - return countLabel(estimate.total, 'attack', 'attacks') + return countLabel(estimate.total, 'planned attack', 'planned attacks') } if (estimate.minimum != null && estimate.maximum != null) { return estimate.minimum === estimate.maximum - ? countLabel(estimate.minimum, 'attack', 'attacks') - : `${formatEstimateValue(estimate.minimum)}-${formatEstimateValue(estimate.maximum)} attacks` + ? countLabel(estimate.minimum, 'planned attack', 'planned attacks') + : `${formatEstimateValue(estimate.minimum)}–${formatEstimateValue(estimate.maximum)} planned attacks` + } + if (estimate.maximum != null) { + return `Up to ${countLabel(estimate.maximum, 'planned attack', 'planned attacks')}` + } + if (estimate.minimum != null) { + return `At least ${countLabel(estimate.minimum, 'planned attack', 'planned attacks')}` + } + return estimate.scope === 'default' + ? 'Select targets to calculate' + : 'Run size is confirmed at launch.' +} + +function baselineCount(estimate: ScenarioRunEstimate): number { + return estimate.components + .filter((component) => component.isBaseline) + .reduce((sum, component) => sum + component.count, 0) +} + +function formatEstimateSummary(estimate: ScenarioRunEstimate): string { + if (estimate.adaptiveDetails) { + const { objectiveCount, techniqueAttemptCountUpperBound } = estimate.adaptiveDetails + const attemptSummary = `up to ${countLabel( + techniqueAttemptCountUpperBound, + 'technique attempt', + 'technique attempts', + )}` + const hasPlannedAttackBound = estimate.total !== null + || estimate.minimum != null + || estimate.maximum != null + return hasPlannedAttackBound + ? `${formatPlannedAttackSummary(estimate)} · ${attemptSummary}` + : `${countLabel(objectiveCount, 'objective', 'objectives')} · ${attemptSummary}` + } + return formatPlannedAttackSummary(estimate) +} + +function operand(id: string, value: string, label: string, result = false, detail?: string): CalculationPart { + return { kind: 'operand', operand: { id, value, label, detail, result } } +} + +function operator(id: string, symbol: CalculationOperator['symbol']): CalculationPart { + return { kind: 'operator', operator: { id, symbol } } +} + +function humanizeLabel(label: string): string { + const words = label.replace(/_/g, ' ').trim() + return words.length > 0 ? `${words[0].toUpperCase()}${words.slice(1)}` : label +} + +function formatDatasetCap(cap: ScenarioRunEstimateDatasetCap): string { + const count = cap.configuredOn === 'dataset' + ? countLabel(cap.count, 'objective', 'objectives') + : formatCount(cap.count) + return `${humanizeLabel(cap.label)}: ${count}` +} + +function semanticFactorLabel(factor: ScenarioRunEstimateFactor): string { + const label = factor.label.toLowerCase() + if (label.includes('seed group') || label === 'objectives') { + return factor.count === 1 ? 'objective' : 'objectives' + } + if (label.includes('technique')) { + return factor.count === 1 ? 'technique' : 'techniques' + } + if (factor.count === 1 && label.endsWith('s')) { + return label.slice(0, -1) + } + return label +} + +function factorPriority(factor: ScenarioRunEstimateFactor): number { + const label = semanticFactorLabel(factor) + if (label === 'technique' || label === 'techniques') return 0 + if (label === 'objective' || label === 'objectives') return 1 + return 2 +} + +function objectiveFactor(component: ScenarioRunEstimateComponent): ScenarioRunEstimateFactor | undefined { + return component.factors.find((factor) => { + const label = semanticFactorLabel(factor) + return label === 'objective' || label === 'objectives' + }) +} + +function resultOperand(estimate: ScenarioRunEstimate): CalculationOperand { + if (estimate.total !== null) { + return { + id: 'result', + value: formatCount(estimate.total), + label: estimate.total === 1 ? 'planned attack' : 'planned attacks', + result: true, + } + } + if (estimate.minimum != null && estimate.maximum != null) { + return { + id: 'result', + value: estimate.minimum === estimate.maximum + ? formatCount(estimate.minimum) + : `${formatCount(estimate.minimum)}–${formatCount(estimate.maximum)}`, + label: estimate.minimum === 1 && estimate.maximum === 1 ? 'planned attack' : 'planned attacks', + result: true, + } } if (estimate.maximum != null) { - return `Up to ${countLabel(estimate.maximum, 'attack', 'attacks')}` + return { + id: 'result', + value: `up to ${formatCount(estimate.maximum)}`, + label: estimate.maximum === 1 ? 'planned attack' : 'planned attacks', + result: true, + } } if (estimate.minimum != null) { - return `At least ${countLabel(estimate.minimum, 'attack', 'attacks')}` + return { + id: 'result', + value: `at least ${formatCount(estimate.minimum)}`, + label: estimate.minimum === 1 ? 'planned attack' : 'planned attacks', + result: true, + } + } + return { id: 'result', value: 'Exact total', label: 'unavailable', result: true } +} + +function adaptivePlannedCalculation(estimate: ScenarioRunEstimate): RunCalculation { + const details = estimate.adaptiveDetails + if (!details) { + throw new Error('Adaptive planned calculation requires adaptive details.') + } + const directBaselineCount = baselineCount(estimate) + const hasExactTotal = estimate.total !== null + || ( + estimate.minimum != null + && estimate.maximum != null + && estimate.minimum === estimate.maximum + ) + const hasPlannedTotal = estimate.total !== null + || estimate.minimum != null + || estimate.maximum != null + const adaptiveAttackCount = hasExactTotal + ? Math.max((estimate.total ?? estimate.maximum ?? 0) - directBaselineCount, 0) + : estimate.maximum != null + ? Math.max(estimate.maximum - directBaselineCount, 0) + : estimate.minimum != null + ? Math.max(estimate.minimum - directBaselineCount, 0) + : details.objectiveCount + const hasAdaptiveRange = directBaselineCount === 0 + && estimate.minimum != null + && estimate.minimum > 0 + && estimate.maximum != null + && estimate.minimum !== estimate.maximum + const adaptiveValue = hasExactTotal + ? formatCount(adaptiveAttackCount) + : hasAdaptiveRange + ? `${formatCount(estimate.minimum ?? 0)}–${formatCount(estimate.maximum ?? 0)}` + : estimate.maximum != null || estimate.minimum == null + ? `up to ${formatCount(adaptiveAttackCount)}` + : `at least ${formatCount(adaptiveAttackCount)}` + const adaptiveLabel = adaptiveAttackCount === 1 ? 'Adaptive attack' : 'Adaptive attacks' + const result = resultOperand(estimate) + const parts: CalculationPart[] = [] + if (directBaselineCount > 0) { + parts.push(operand( + 'baseline', + formatCount(directBaselineCount), + directBaselineCount === 1 ? 'direct baseline attack' : 'direct baseline attacks', + )) + parts.push(operator('baseline-plus', '+')) + } + parts.push(operand('adaptive-attacks', adaptiveValue, adaptiveLabel)) + if (hasPlannedTotal) { + parts.push(operator('planned-equals', '=')) + parts.push({ kind: 'operand', operand: result }) + } + + const adaptivePhrase = `${adaptiveValue} ${adaptiveLabel}` + const resultPhrase = `${result.value} ${result.label}` + const plannedResultPhrase = hasPlannedTotal + ? ` equals ${resultPhrase}.` + : '. Planned total is confirmed at launch.' + const accessibleLabel = directBaselineCount > 0 + ? `Direct baseline comparison is included: ${countLabel( + directBaselineCount, + 'direct baseline attack', + 'direct baseline attacks', + )} plus ${adaptivePhrase}${plannedResultPhrase}` + : `Direct baseline comparison is not included: ${adaptivePhrase}${plannedResultPhrase}` + return { parts, accessibleLabel } +} + +function adaptiveWorkCalculation(estimate: ScenarioRunEstimate): RunCalculation { + const details = estimate.adaptiveDetails + if (!details) { + throw new Error('Adaptive work calculation requires adaptive details.') + } + const objectiveLabel = details.objectiveCount === 1 ? 'objective' : 'objectives' + const techniqueLabel = details.techniquesPerObjectiveUpperBound === 1 + ? 'technique per objective' + : 'techniques per objective' + const attemptLabel = details.techniqueAttemptCountUpperBound === 1 + ? 'technique attempt' + : 'technique attempts' + const capProvenance = { + selectedCandidateCount: details.selectedCandidateTechniqueCount, + compatibleCandidateCount: details.candidateTechniqueCount, + limit: details.maxAttemptsPerObjective, + effectiveMaximum: details.techniquesPerObjectiveUpperBound, + } + const effectiveCapRule = formatAdaptiveCapMetadata(capProvenance) + const accessibleCapRule = formatAdaptiveCapAccessibleRule(capProvenance) + return { + parts: [ + operand('adaptive-objectives', formatCount(details.objectiveCount), objectiveLabel), + operator('adaptive-multiply', '×'), + operand( + 'adaptive-techniques', + `up to ${formatCount(details.techniquesPerObjectiveUpperBound)}`, + techniqueLabel, + false, + effectiveCapRule, + ), + operator('adaptive-equals', '='), + operand( + 'adaptive-result', + `up to ${formatCount(details.techniqueAttemptCountUpperBound)}`, + attemptLabel, + true, + ), + ], + accessibleLabel: `${countLabel(details.objectiveCount, 'objective', 'objectives')} multiplied by up to ${ + countLabel(details.techniquesPerObjectiveUpperBound, 'technique per objective', 'techniques per objective') + }, ${accessibleCapRule}, equals up to ${ + countLabel(details.techniqueAttemptCountUpperBound, 'technique attempt', 'technique attempts') + }.`, } - return 'Attack count varies' } -function formatComponentFormula(component: ScenarioRunEstimateComponent): string { - return `${component.label}: ${formatEstimateValue(component.count)}` +function adaptiveWorkContext(estimate: ScenarioRunEstimate): string { + const details = estimate.adaptiveDetails + if (!details) { + throw new Error('Adaptive work context requires adaptive details.') + } + const compatibilityContext = details.compatibilityMayReduceAttempts + ? ' Compatibility may reduce how many candidates each objective can try.' + : '' + return `Technique-attempt totals exclude multi-turn target exchanges and retries. Adaptive stops each objective after the first successful technique.${compatibilityContext}` +} + +function homogeneousTechniqueCalculation( + components: ScenarioRunEstimateComponent[], +): CalculationPart[] | null { + if (components.length < 2 || components.some((component) => component.condition !== null)) { + return null + } + const objectiveCounts = components.map((component) => objectiveFactor(component)?.count) + if (objectiveCounts.some((count) => count === undefined)) { + return null + } + const firstCount = objectiveCounts[0] + if (!objectiveCounts.every((count) => count === firstCount)) { + return null + } + return [ + operand( + 'technique-count', + formatCount(components.length), + components.length === 1 ? 'technique' : 'techniques', + ), + operator('technique-multiply', '×'), + operand( + 'objective-count', + formatCount(firstCount ?? 0), + firstCount === 1 ? 'objective' : 'objectives', + ), + ] } -function formatBackendFormula(estimate: ScenarioRunEstimate): string { - const components = estimate.components.length > 0 - ? estimate.components.map(formatComponentFormula).join(' + ') - : 'No additive components supplied' - const total = estimate.total === null - ? 'conditional total' - : formatEstimateValue(estimate.total) - return `${components} = ${total}` +function componentTerms(components: ScenarioRunEstimateComponent[]): CalculationPart[] { + const homogeneous = homogeneousTechniqueCalculation(components) + if (homogeneous) { + return homogeneous + } + if (components.length === 1 && components[0].condition === null && components[0].factors.length > 0) { + return [...components[0].factors] + .sort((left, right) => factorPriority(left) - factorPriority(right)) + .flatMap((factor, index) => [ + ...(index > 0 ? [operator(`factor-${index}-multiply`, '×')] : []), + operand(`factor-${factor.id}`, formatCount(factor.count), semanticFactorLabel(factor)), + ]) + } + return components.flatMap((component, index) => { + const objectiveCount = objectiveFactor(component)?.count + const value = formatCount(objectiveCount ?? component.count) + const unit = objectiveCount === undefined + ? component.count === 1 ? 'planned attack' : 'planned attacks' + : objectiveCount === 1 ? 'objective' : 'objectives' + const condition = component.condition ? ' · if supported' : '' + return [ + ...(index > 0 ? [operator(`component-${index}-plus`, '+')] : []), + operand( + `component-${component.id}`, + value, + `${unit} · ${humanizeLabel(component.label)}${condition}`, + ), + ] + }) +} + +function ordinaryCalculation(estimate: ScenarioRunEstimate): RunCalculation { + const baselineCount = estimate.components + .filter((component) => component.isBaseline) + .reduce((sum, component) => sum + component.count, 0) + const attackComponents = estimate.components.filter((component) => !component.isBaseline) + const resultOnly = baselineCount === 0 + && attackComponents.length === 1 + && attackComponents[0].factors.length === 0 + && attackComponents[0].condition === null + const attackParts = resultOnly ? [] : componentTerms(attackComponents) + const hasMultiplication = attackParts.some( + (part) => part.kind === 'operator' && part.operator.symbol === '×', + ) + const parts: CalculationPart[] = [] + if (baselineCount > 0 && hasMultiplication) { + parts.push(operator('attack-open', '(')) + } + parts.push(...attackParts) + if (baselineCount > 0 && hasMultiplication) { + parts.push(operator('attack-close', ')')) + } + if (baselineCount > 0) { + if (attackParts.length > 0) { + parts.push(operator('baseline-plus', '+')) + } + parts.push(operand( + 'baseline', + formatCount(baselineCount), + baselineCount === 1 ? 'direct baseline attack' : 'direct baseline attacks', + )) + } + const result = resultOperand(estimate) + if (parts.length > 0) { + parts.push(operator('total-equals', '=')) + } + parts.push({ kind: 'operand', operand: result }) + + const visibleExpression = parts.map((part) => part.kind === 'operator' + ? part.operator.symbol + : `${part.operand.value} ${part.operand.label}`).join(' ') + return { + parts, + accessibleLabel: `${visibleExpression + .replace(/×/g, 'multiplied by') + .replace(/\+/g, 'plus') + .replace(/=/g, 'equals')}.`, + context: estimate.total === null && estimate.minimum == null && estimate.maximum == null + ? formatEstimateSummary(estimate) + : undefined, + } } -export function ScenarioRunEstimateSummary({ state, compact = false }: ScenarioRunEstimateSummaryProps) { +export function ScenarioRunEstimateSummary({ state }: ScenarioRunEstimateSummaryProps) { const styles = useScenarioRunEstimateStyles() const estimate = stateEstimate(state) return (
- {!compact && {statusLabel(state)}} + {statusLabel(state)} {estimate && ( - {formatPlannedAttackSummary(estimate)} + {formatEstimateSummary(estimate)} )} - {compact && !estimate && {statusLabel(state)}}
- {!compact && {scopeLabel(state)}}
) } -export function ScenarioRunEstimateDetails({ state }: ScenarioRunEstimateDetailsProps) { +function RunCalculationView({ + calculation, + heading, + idPrefix, + testId, +}: { + calculation: RunCalculation + heading: string + idPrefix: string + testId: string +}) { + const styles = useScenarioRunEstimateStyles() + const headingId = `${idPrefix}-calculation` + + return ( +
+ {heading} +
+ {calculation.parts.map((part) => part.kind === 'operator' ? ( + + ) : ( + + ))} +
+ {calculation.summary && ( + + {calculation.summary} + + )} + {calculation.context && ( + {calculation.context} + )} +
+ ) +} + +function EstimateSources({ estimate }: { estimate: ScenarioRunEstimate }) { + const styles = useScenarioRunEstimateStyles() + if (estimate.datasets.length === 0) { + return null + } + const { commonCaps, residualCapsByDatasetId } = normalizeDatasetCaps(estimate.datasets) + + return ( +
+ {commonCaps.length > 0 && ( + + {commonCaps.map(formatDatasetCap).join(' · ')} + + )} + {estimate.datasets.map((dataset) => { + const residualCaps = residualCapsByDatasetId.get(dataset.id) ?? [] + return ( +
+ + {countLabel(dataset.selectedSeedGroupCount, 'objective', 'objectives')} from {dataset.name} + {dataset.logicalSeedGroupCount !== dataset.selectedSeedGroupCount + ? ` · ${formatCount(dataset.logicalSeedGroupCount)} available` + : ''} + + {dataset.selectionNote && ( + {dataset.selectionNote} + )} + {residualCaps.length > 0 && ( + + {residualCaps.map(formatDatasetCap).join(' · ')} + + )} +
+ ) + })} +
+ ) +} + +export function ScenarioRunEstimateDetails({ + state, + idPrefix = 'scenario-run-estimate', +}: ScenarioRunEstimateDetailsProps) { const styles = useScenarioRunEstimateStyles() if (state.status === 'loading') { return (
- + + {scopeLabel(state)}
) } @@ -140,19 +611,46 @@ export function ScenarioRunEstimateDetails({ state }: ScenarioRunEstimateDetails if (state.status === 'unavailable') { return (
- {state.label} + + {state.label} {state.note && {state.note}}
) } const { estimate } = state + const hasAdaptiveDetails = estimate.adaptiveDetails !== null return (
- - {formatPlannedAttackSummary(estimate)} + {hasAdaptiveDetails ? ( + <> + + + + ) : ( + + )} + + + {hasAdaptiveDetails + ? adaptiveWorkContext(estimate) + : `Retries are ${estimate.retriesIncluded ? 'included' : 'not included'}.`} - {formatBackendFormula(estimate)}
) } diff --git a/frontend/src/components/Scenarios/scenarioAdaptiveCap.test.ts b/frontend/src/components/Scenarios/scenarioAdaptiveCap.test.ts new file mode 100644 index 0000000000..1879453e53 --- /dev/null +++ b/frontend/src/components/Scenarios/scenarioAdaptiveCap.test.ts @@ -0,0 +1,42 @@ +import { + formatAdaptiveCapAccessibleRule, + formatAdaptiveCapFeedback, + formatAdaptiveCapMetadata, +} from './scenarioAdaptiveCap' + +describe('scenarioAdaptiveCap', () => { + it.each([ + [1, 1, '1 selected candidate · limit 3'], + [4, 4, '4 selected candidates · limit 3'], + [2, 1, '1 compatible candidate from 2 selected · limit 3'], + [4, 2, '2 compatible candidates from 4 selected · limit 3'], + ])( + 'formats metadata for %i selected and %i compatible candidates', + (selectedCandidateCount, compatibleCandidateCount, expected) => { + expect(formatAdaptiveCapMetadata({ + selectedCandidateCount, + compatibleCandidateCount, + limit: 3, + effectiveMaximum: 2, + })).toBe(expected) + }, + ) + + it('formats feedback with the effective maximum', () => { + expect(formatAdaptiveCapFeedback({ + selectedCandidateCount: 4, + compatibleCandidateCount: 2, + limit: 3, + effectiveMaximum: 2, + })).toBe('2 compatible candidates from 4 selected · limit 3 · effective maximum 2.') + }) + + it('formats the accessible minimum rule', () => { + expect(formatAdaptiveCapAccessibleRule({ + selectedCandidateCount: 1, + compatibleCandidateCount: 1, + limit: 3, + effectiveMaximum: 1, + })).toBe('the smaller of 1 selected candidate and limit 3') + }) +}) diff --git a/frontend/src/components/Scenarios/scenarioAdaptiveCap.ts b/frontend/src/components/Scenarios/scenarioAdaptiveCap.ts new file mode 100644 index 0000000000..6e7f4a3b26 --- /dev/null +++ b/frontend/src/components/Scenarios/scenarioAdaptiveCap.ts @@ -0,0 +1,33 @@ +interface AdaptiveCapProvenance { + selectedCandidateCount: number + compatibleCandidateCount: number + limit: number + effectiveMaximum: number +} + +function candidateContext({ + selectedCandidateCount, + compatibleCandidateCount, +}: Pick): string { + const compatibleLabel = compatibleCandidateCount === 1 ? 'compatible candidate' : 'compatible candidates' + const selectedLabel = selectedCandidateCount === 1 ? 'selected candidate' : 'selected candidates' + return compatibleCandidateCount < selectedCandidateCount + ? `${compatibleCandidateCount.toLocaleString()} ${compatibleLabel} from ${ + selectedCandidateCount.toLocaleString() + } selected` + : `${selectedCandidateCount.toLocaleString()} ${selectedLabel}` +} + +export function formatAdaptiveCapMetadata(provenance: AdaptiveCapProvenance): string { + return `${candidateContext(provenance)} · limit ${provenance.limit.toLocaleString()}` +} + +export function formatAdaptiveCapFeedback(provenance: AdaptiveCapProvenance): string { + return `${formatAdaptiveCapMetadata(provenance)} · effective maximum ${ + provenance.effectiveMaximum.toLocaleString() + }.` +} + +export function formatAdaptiveCapAccessibleRule(provenance: AdaptiveCapProvenance): string { + return `the smaller of ${candidateContext(provenance)} and limit ${provenance.limit.toLocaleString()}` +} diff --git a/frontend/src/components/Scenarios/scenarioDatasetCaps.test.ts b/frontend/src/components/Scenarios/scenarioDatasetCaps.test.ts new file mode 100644 index 0000000000..b3c0e2d9a9 --- /dev/null +++ b/frontend/src/components/Scenarios/scenarioDatasetCaps.test.ts @@ -0,0 +1,70 @@ +import type { ScenarioRunEstimateDataset, ScenarioRunEstimateDatasetCap } from '@/types' + +import { normalizeDatasetCaps } from './scenarioDatasetCaps' + +const SHARED_CAP: ScenarioRunEstimateDatasetCap = { + id: 'shared-cap', + label: 'combined cap', + count: 10, + configuredOn: 'compound', + datasetName: null, +} + +function makeDataset( + name: string, + configuredCaps: ScenarioRunEstimateDatasetCap[], +): ScenarioRunEstimateDataset { + return { + id: name, + name, + kind: 'dataset', + logicalSeedGroupCount: 20, + selectedSeedGroupCount: 5, + configuredCaps, + selectionNote: null, + } +} + +describe('scenarioDatasetCaps', () => { + it('returns empty normalized collections without datasets', () => { + const normalized = normalizeDatasetCaps([]) + + expect(normalized.commonCaps).toEqual([]) + expect(normalized.residualCapsByDatasetId.size).toBe(0) + }) + + it('lifts compound caps and preserves single-dataset row caps', () => { + const rowCap: ScenarioRunEstimateDatasetCap = { + ...SHARED_CAP, + id: 'row-cap', + configuredOn: 'dataset', + datasetName: 'alpha', + } + const dataset = makeDataset('alpha', [SHARED_CAP, rowCap]) + const normalized = normalizeDatasetCaps([dataset]) + + expect(normalized.commonCaps).toEqual([SHARED_CAP]) + expect(normalized.residualCapsByDatasetId.get('alpha')).toEqual([rowCap]) + }) + + it('preserves cap multiplicity while separating universal and residual row caps', () => { + const rowCap: ScenarioRunEstimateDatasetCap = { + id: 'row-cap', + label: 'per-dataset cap', + count: 5, + configuredOn: 'dataset', + datasetName: null, + } + const datasets = [ + makeDataset('alpha', [SHARED_CAP, rowCap, { ...rowCap, id: 'row-cap-duplicate' }]), + makeDataset('beta', [{ ...SHARED_CAP, id: 'shared-cap-beta' }, rowCap]), + ] + const normalized = normalizeDatasetCaps(datasets) + + expect(normalized.commonCaps).toEqual([SHARED_CAP, rowCap]) + expect(normalized.residualCapsByDatasetId.get('alpha')).toEqual([ + { ...rowCap, id: 'row-cap-duplicate' }, + ]) + expect(normalized.residualCapsByDatasetId.get('beta')).toEqual([]) + }) +}) diff --git a/frontend/src/components/Scenarios/scenarioDatasetCaps.ts b/frontend/src/components/Scenarios/scenarioDatasetCaps.ts new file mode 100644 index 0000000000..9c1b86c6f1 --- /dev/null +++ b/frontend/src/components/Scenarios/scenarioDatasetCaps.ts @@ -0,0 +1,87 @@ +import type { + ScenarioRunEstimateDataset, + ScenarioRunEstimateDatasetCap, +} from '@/types' + +interface NormalizedDatasetCaps { + readonly commonCaps: ScenarioRunEstimateDatasetCap[] + readonly residualCapsByDatasetId: ReadonlyMap +} + +function semanticCapKey(cap: ScenarioRunEstimateDatasetCap): string { + return JSON.stringify([cap.label, cap.count, cap.configuredOn]) +} + +function capOccurrences(caps: ScenarioRunEstimateDatasetCap[]): Map { + const occurrences = new Map() + for (const cap of caps) { + const key = semanticCapKey(cap) + occurrences.set(key, (occurrences.get(key) ?? 0) + 1) + } + return occurrences +} + +export function normalizeDatasetCaps(datasets: ScenarioRunEstimateDataset[]): NormalizedDatasetCaps { + const commonCaps: ScenarioRunEstimateDatasetCap[] = [] + const seenListLevelCaps = new Set() + + for (const dataset of datasets) { + for (const cap of dataset.configuredCaps) { + if (cap.configuredOn !== 'compound') { + continue + } + const key = semanticCapKey(cap) + if (!seenListLevelCaps.has(key)) { + seenListLevelCaps.add(key) + commonCaps.push(cap) + } + } + } + + const universalRowCapOccurrences = datasets.length > 1 + ? capOccurrences(datasets[0].configuredCaps.filter((cap) => cap.configuredOn !== 'compound')) + : new Map() + for (const dataset of datasets.slice(1)) { + const datasetOccurrences = capOccurrences( + dataset.configuredCaps.filter((cap) => cap.configuredOn !== 'compound'), + ) + for (const [key, count] of universalRowCapOccurrences) { + universalRowCapOccurrences.set(key, Math.min(count, datasetOccurrences.get(key) ?? 0)) + } + } + + const emittedRowCapOccurrences = new Map() + if (datasets.length > 0) { + for (const cap of datasets[0].configuredCaps) { + if (cap.configuredOn === 'compound') { + continue + } + const key = semanticCapKey(cap) + const emitted = emittedRowCapOccurrences.get(key) ?? 0 + if (emitted < (universalRowCapOccurrences.get(key) ?? 0)) { + commonCaps.push(cap) + emittedRowCapOccurrences.set(key, emitted + 1) + } + } + } + + const residualCapsByDatasetId = new Map() + for (const dataset of datasets) { + const consumedRowCapOccurrences = new Map() + const residualCaps = dataset.configuredCaps.filter((cap) => { + if (cap.configuredOn === 'compound') { + return false + } + const key = semanticCapKey(cap) + const consumed = consumedRowCapOccurrences.get(key) ?? 0 + if (consumed < (universalRowCapOccurrences.get(key) ?? 0)) { + consumedRowCapOccurrences.set(key, consumed + 1) + return false + } + return true + }) + residualCapsByDatasetId.set(dataset.id, residualCaps) + } + + return { commonCaps, residualCapsByDatasetId } +} diff --git a/frontend/src/components/Scenarios/scenarioRunEstimateAdapter.ts b/frontend/src/components/Scenarios/scenarioRunEstimateAdapter.ts index 4fd7cbcb14..63f3bed303 100644 --- a/frontend/src/components/Scenarios/scenarioRunEstimateAdapter.ts +++ b/frontend/src/components/Scenarios/scenarioRunEstimateAdapter.ts @@ -1,7 +1,10 @@ import type { + ScenarioDefaultRunSizeEstimate, ScenarioRunEstimate, + ScenarioRunEstimateAdaptiveDetails, ScenarioRunEstimateDataset, ScenarioRunEstimateDatasetCap, + ScenarioRunEstimateFactor, ScenarioRunEstimateResult, ScenarioRunSizeEstimateResponse, } from '@/types' @@ -14,7 +17,7 @@ function nextStableId(prefix: string, label: string, occurrences: Map() return caps.map((cap) => ({ @@ -27,7 +30,7 @@ function mapDatasetCaps( } function mapDatasets( - datasets: ScenarioRunSizeEstimateResponse['datasets'], + datasets: ScenarioDefaultRunSizeEstimate['datasets'], ): ScenarioRunEstimateDataset[] { const occurrences = new Map() return datasets.map((dataset) => { @@ -44,16 +47,47 @@ function mapDatasets( }) } +function mapFactors( + componentId: string, + factors: NonNullable, +): ScenarioRunEstimateFactor[] { + const occurrences = new Map() + return factors.map((factor) => ({ + id: nextStableId(`${componentId}:factor`, factor.label, occurrences), + label: factor.label, + count: factor.count, + })) +} + +function mapAdaptiveDetails( + details: NonNullable, +): ScenarioRunEstimateAdaptiveDetails { + return { + objectiveCount: details.objective_count, + selectedCandidateTechniqueCount: + details.selected_candidate_technique_count ?? details.candidate_technique_count, + candidateTechniqueCount: details.candidate_technique_count, + maxAttemptsPerObjective: details.max_attempts_per_objective, + techniquesPerObjectiveUpperBound: details.techniques_per_objective_upper_bound, + techniqueAttemptCountUpperBound: details.technique_attempt_count_upper_bound, + stopOnFirstSuccess: details.stop_on_first_success, + compatibilityMayReduceAttempts: details.compatibility_may_reduce_attempts, + } +} + export function mapScenarioRunEstimate( - response: ScenarioRunSizeEstimateResponse, + response: ScenarioDefaultRunSizeEstimate | ScenarioRunSizeEstimateResponse, scope: ScenarioRunEstimate['scope'], ): ScenarioRunEstimateResult { - if ( - response.estimated_attack_count === null - && response.minimum_attack_count == null - && response.maximum_attack_count == null - && response.components.length === 0 - ) { + const isRichEstimate = 'status' in response + const total = isRichEstimate ? response.total_attack_count : response.estimated_attack_count + const unavailable = isRichEstimate + ? response.status === 'unavailable' + : total === null + && response.minimum_attack_count == null + && response.maximum_attack_count == null + && response.components.length === 0 + if (unavailable) { return { status: 'unavailable', scope, @@ -66,26 +100,34 @@ export function mapScenarioRunEstimate( const componentOccurrences = new Map() const estimate: ScenarioRunEstimate = { + version: isRichEstimate ? response.version : 1, scope, - total: response.estimated_attack_count, + total, minimum: response.minimum_attack_count ?? null, maximum: response.maximum_attack_count ?? null, + condition: isRichEstimate ? response.condition : null, components: response.components.map((component) => { const id = nextStableId('component', component.label, componentOccurrences) return { id, label: component.label, count: component.count, + factors: mapFactors(id, component.factors ?? []), isBaseline: component.is_baseline, + condition: component.condition, note: component.note, } }), datasets: mapDatasets(response.datasets), + adaptiveDetails: isRichEstimate && response.adaptive_details + ? mapAdaptiveDetails(response.adaptive_details) + : null, effectiveParameters: response.effective_parameters ?? {}, note: response.note, + retriesIncluded: isRichEstimate ? response.retries_included : false, } - return response.estimated_attack_count === null + return (isRichEstimate ? response.status === 'conditional' : total === null) ? { status: 'conditional', estimate } : { status: 'available', estimate } } diff --git a/frontend/src/components/Scenarios/scenarioTechniqueSets.test.ts b/frontend/src/components/Scenarios/scenarioTechniqueSets.test.ts new file mode 100644 index 0000000000..dd604b3fd0 --- /dev/null +++ b/frontend/src/components/Scenarios/scenarioTechniqueSets.test.ts @@ -0,0 +1,73 @@ +import type { RegisteredScenario } from '@/types' + +import { + techniqueSetDisplayName, + techniqueSetMembers, + techniqueSetName, + techniqueSetOptionLabel, +} from './scenarioTechniqueSets' + +function makeScenario(overrides: Partial = {}): RegisteredScenario { + return { + scenario_name: 'test.scenario', + scenario_type: 'TestScenario', + scenario_version: 1, + description: 'Test scenario.', + description_markdown: 'Test scenario.', + default_technique: 'default', + default_techniques: ['crescendo'], + aggregate_techniques: ['default', 'quick_set'], + aggregate_technique_expansions: { + quick_set: ['crescendo', 'crescendo', 'pair'], + }, + all_techniques: ['crescendo', 'pair'], + default_datasets: [], + dataset_size_limit: { + default_scope: 'none', + default_count: null, + override_scope: 'unsupported', + }, + baseline_policy: 'forbidden', + include_baseline_by_default: false, + supported_parameters: [], + default_run_size: { + version: 1, + status: 'unavailable', + total_attack_count: null, + minimum_attack_count: null, + maximum_attack_count: null, + condition: null, + components: [], + datasets: [], + adaptive_details: null, + note: null, + retries_included: false, + }, + ...overrides, + } +} + +describe('scenarioTechniqueSets', () => { + it('formats known, custom, and empty technique-set names', () => { + expect(techniqueSetName('default')).toBe('Recommended') + expect(techniqueSetName('custom_red_team')).toBe('Custom red team') + expect(techniqueSetName('')).toBe('') + }) + + it('expands named sets, removes duplicates, and falls back to default members', () => { + const scenario = makeScenario() + + expect(techniqueSetMembers(scenario, 'quick_set')).toEqual(['crescendo', 'pair']) + expect(techniqueSetMembers(scenario, 'default')).toEqual(['crescendo']) + expect(techniqueSetMembers(scenario, 'unknown_set')).toEqual([]) + }) + + it('labels default and custom sets with singular and plural member counts', () => { + const scenario = makeScenario() + + expect(techniqueSetDisplayName(scenario, 'default')).toBe('Recommended (default)') + expect(techniqueSetDisplayName(scenario, 'quick_set')).toBe('Quick set') + expect(techniqueSetOptionLabel(scenario, 'default')).toBe('Recommended (default) — 1 technique') + expect(techniqueSetOptionLabel(scenario, 'quick_set')).toBe('Quick set (2 techniques)') + }) +}) diff --git a/frontend/src/services/api.test.ts b/frontend/src/services/api.test.ts index 9de3b57847..68690f5ba5 100644 --- a/frontend/src/services/api.test.ts +++ b/frontend/src/services/api.test.ts @@ -21,6 +21,7 @@ import { configurationApi, targetsApi, attacksApi, + datasetsApi, scenariosApi, } from "./api"; @@ -564,6 +565,22 @@ describe("api service", () => { }); }); + describe("datasetsApi", () => { + it("lists registered datasets", async () => { + const mockResponse = { + data: { + items: [{ name: "harmbench" }, { name: "xstest" }], + }, + }; + (apiClient.get as jest.Mock).mockResolvedValueOnce(mockResponse); + + const result = await datasetsApi.listDatasets(); + + expect(apiClient.get).toHaveBeenCalledWith("/datasets"); + expect(result).toEqual(mockResponse.data); + }); + }); + describe("scenariosApi", () => { it("lists the scenario catalog with default params", async () => { const mockResponse = { @@ -647,9 +664,15 @@ describe("api service", () => { it("posts the exact estimate request and forwards cancellation", async () => { const mockResponse = { data: { - estimated_attack_count: 8, + version: 1, + status: "exact", + total_attack_count: 8, + minimum_attack_count: null, + maximum_attack_count: null, + condition: null, components: [], datasets: [], + adaptive_details: null, note: null, }, }; @@ -676,7 +699,56 @@ describe("api service", () => { request, { signal: controller.signal } ); - expect(result.estimated_attack_count).toBe(8); + expect(result.total_attack_count).toBe(8); + }); + + it("preserves Adaptive conditional work metadata in the initial estimate response", async () => { + const mockResponse = { + data: { + version: 1, + status: "conditional", + total_attack_count: null, + minimum_attack_count: null, + maximum_attack_count: null, + condition: null, + components: [], + datasets: [], + adaptive_details: { + objective_count: 21, + selected_candidate_technique_count: 2, + candidate_technique_count: 2, + max_attempts_per_objective: 3, + techniques_per_objective_upper_bound: 2, + technique_attempt_count_upper_bound: 42, + stop_on_first_success: true, + compatibility_may_reduce_attempts: true, + }, + note: null, + retries_included: false, + }, + }; + (apiClient.post as jest.Mock).mockResolvedValueOnce(mockResponse); + const controller = new AbortController(); + const request = { + target_name: "target-a", + techniques: ["default"], + include_baseline: true, + scenario_params: { max_attempts_per_objective: 3 }, + }; + + const result = await scenariosApi.estimateRun( + "adaptive.text_adaptive", + request, + controller.signal + ); + + expect(apiClient.post).toHaveBeenCalledWith( + "/scenarios/catalog/adaptive.text_adaptive/estimate", + request, + { signal: controller.signal } + ); + expect(result.adaptive_details).toEqual(mockResponse.data.adaptive_details); + expect(result.total_attack_count).toBeNull(); }); it("posts the exact RunScenarioRequest payload to start a run", async () => { diff --git a/frontend/src/services/api.ts b/frontend/src/services/api.ts index 6cf66da4c6..5a3544f44f 100644 --- a/frontend/src/services/api.ts +++ b/frontend/src/services/api.ts @@ -30,10 +30,11 @@ import type { CreateConversationRequest, CreateConversationResponse, ChangeMainConversationResponse, + DatasetListResponse, ListRegisteredScenariosResponse, RegisteredScenario, RunScenarioRequest, - ScenarioRunSizeEstimateResponse, + ScenarioDefaultRunSizeEstimate, ScenarioRunSizeEstimateRequest, ScenarioRunSummary, ScenarioRunListResponse, @@ -409,6 +410,13 @@ export const labelsApi = { }, } +export const datasetsApi = { + listDatasets: async (): Promise => { + const response = await apiClient.get('/datasets') + return response.data + }, +} + export const scenariosApi = { /** * Lists one page of the scenario catalog. Callers that need the full @@ -445,7 +453,7 @@ export const scenariosApi = { scenarioName: string, request: ScenarioRunSizeEstimateRequest, signal?: AbortSignal, - ): Promise => { + ): Promise => { const response = await apiClient.post( `/scenarios/catalog/${encodeURIComponent(scenarioName)}/estimate`, request, diff --git a/frontend/src/types/index.ts b/frontend/src/types/index.ts index da3a6e668e..cf37efb790 100644 --- a/frontend/src/types/index.ts +++ b/frontend/src/types/index.ts @@ -484,6 +484,16 @@ export interface ChangeMainConversationResponse { conversation_id: string } +// --- Datasets --- + +export interface DatasetInfo { + name: string +} + +export interface DatasetListResponse { + items: DatasetInfo[] +} + // --- Scenarios --- export interface RegisteredScenario { @@ -499,10 +509,11 @@ export interface RegisteredScenario { all_techniques: string[] technique_summaries: ScenarioTechniqueSummary[] default_datasets: string[] + dataset_size_limit: ScenarioDatasetSizeLimit baseline_policy: 'enabled' | 'disabled' | 'forbidden' include_baseline_by_default: boolean supported_parameters: Parameter[] - default_run_size: ScenarioRunSizeEstimateResponse + default_run_size: ScenarioDefaultRunSizeEstimate | ScenarioRunSizeEstimateResponse } export interface ScenarioTechniqueSummary { @@ -533,13 +544,33 @@ export interface RunScenarioRequest { scenario_result_id?: string | null } +export type ScenarioRunSizeEstimateStatus = 'exact' | 'conditional' | 'unavailable' + +export interface ScenarioRunSizeFactor { + label: string + count: number +} + export interface ScenarioRunSizeComponent { label: string count: number + factors?: ScenarioRunSizeFactor[] is_baseline: boolean + condition?: 'target_capabilities' | 'launch_configuration' | null note: string | null } +export interface ScenarioAdaptiveRunSizeDetails { + objective_count: number + selected_candidate_technique_count?: number + candidate_technique_count: number + max_attempts_per_objective: number + techniques_per_objective_upper_bound: number + technique_attempt_count_upper_bound: number + stop_on_first_success: true + compatibility_may_reduce_attempts: true +} + export interface ScenarioDatasetSizeCap { label: string count: number @@ -556,6 +587,12 @@ export interface ScenarioDatasetSummary { selection_note: string | null } +export interface ScenarioDatasetSizeLimit { + default_scope: 'none' | 'per_dataset' | 'combined' | 'heterogeneous' + default_count: number | null + override_scope: 'per_dataset' | 'combined' | 'unsupported' +} + export interface ScenarioRunSizeEstimateResponse { estimated_attack_count: number | null minimum_attack_count?: number | null @@ -566,6 +603,21 @@ export interface ScenarioRunSizeEstimateResponse { note: string | null } +export interface ScenarioDefaultRunSizeEstimate { + version: 1 + status: ScenarioRunSizeEstimateStatus + total_attack_count: number | null + minimum_attack_count?: number | null + maximum_attack_count?: number | null + condition?: 'target_capabilities' | 'launch_configuration' | null + components: ScenarioRunSizeComponent[] + datasets: ScenarioDatasetSummary[] + adaptive_details?: ScenarioAdaptiveRunSizeDetails | null + effective_parameters?: Record + note: string | null + retries_included: false +} + export interface ScenarioRunSizeEstimateRequest { target_name?: string | null techniques?: string[] | null @@ -580,10 +632,29 @@ export interface ScenarioRunEstimateComponent { id: string label: string count: number + factors: ScenarioRunEstimateFactor[] isBaseline: boolean + condition?: 'target_capabilities' | 'launch_configuration' | null note: string | null } +export interface ScenarioRunEstimateFactor { + id: string + label: string + count: number +} + +export interface ScenarioRunEstimateAdaptiveDetails { + objectiveCount: number + selectedCandidateTechniqueCount: number + candidateTechniqueCount: number + maxAttemptsPerObjective: number + techniquesPerObjectiveUpperBound: number + techniqueAttemptCountUpperBound: number + stopOnFirstSuccess: true + compatibilityMayReduceAttempts: true +} + export interface ScenarioRunEstimateDatasetCap { id: string label: string @@ -603,14 +674,18 @@ export interface ScenarioRunEstimateDataset { } export interface ScenarioRunEstimate { + version: number scope: 'default' | 'request' total: number | null minimum?: number | null maximum?: number | null + condition?: 'target_capabilities' | 'launch_configuration' | null components: ScenarioRunEstimateComponent[] datasets: ScenarioRunEstimateDataset[] + adaptiveDetails?: ScenarioRunEstimateAdaptiveDetails | null effectiveParameters: Record note: string | null + retriesIncluded: boolean } export type ScenarioRunEstimateResult = @@ -651,7 +726,7 @@ export type ScenarioRunEstimator = ( scenarioName: string, request: ScenarioRunSizeEstimateRequest, signal?: AbortSignal, -) => Promise +) => Promise export interface AttackErrorSummary { atomic_attack_name: string diff --git a/pyrit/analytics/technique_analysis.py b/pyrit/analytics/technique_analysis.py index b892946113..a804a510ac 100644 --- a/pyrit/analytics/technique_analysis.py +++ b/pyrit/analytics/technique_analysis.py @@ -5,6 +5,7 @@ from __future__ import annotations +from collections import Counter, defaultdict from typing import TYPE_CHECKING from pyrit.analytics.result_analysis import AttackStats, _compute_stats @@ -12,11 +13,33 @@ from pyrit.models import AttackOutcome if TYPE_CHECKING: - from collections.abc import Sequence + from collections.abc import Iterable, Mapping, Sequence from pyrit.memory.memory_interface import MemoryInterface +def _compute_grouped_outcome_stats(grouped_outcomes: Iterable[tuple[str, AttackOutcome]]) -> dict[str, AttackStats]: + """ + Aggregate keyed outcomes into attack statistics. + + Returns: + dict[str, AttackStats]: Statistics keyed by the caller's grouping value. + """ + counts: dict[str, Counter[AttackOutcome]] = defaultdict(Counter) + for key, outcome in grouped_outcomes: + counts[key][outcome] += 1 + + return { + key: _compute_stats( + successes=outcomes[AttackOutcome.SUCCESS], + failures=outcomes[AttackOutcome.FAILURE], + undetermined=outcomes[AttackOutcome.UNDETERMINED], + errors=outcomes[AttackOutcome.ERROR], + ) + for key, outcomes in counts.items() + } + + def compute_technique_stats( *, technique_eval_hashes: Sequence[str], @@ -61,24 +84,88 @@ def compute_technique_stats( ) requested = set(technique_eval_hashes) - counts: dict[str, tuple[int, int, int, int]] = {} + grouped_outcomes: list[tuple[str, AttackOutcome]] = [] for result in results: identifier = result.atomic_attack_identifier eval_hash = identifier.eval_hash if identifier is not None else None if eval_hash is None or eval_hash not in requested: continue + grouped_outcomes.append((eval_hash, result.outcome)) + + return _compute_grouped_outcome_stats(grouped_outcomes) + + +def compute_labeled_technique_stats( + *, + technique_identifiers: Sequence[str], + label_name: str, + technique_eval_hashes_by_identifier: Mapping[str, str] | None = None, + scenario_result_id: str | None = None, + targeted_harm_categories: Sequence[str] | None = None, + memory: MemoryInterface | None = None, +) -> dict[str, AttackStats]: + """ + Compute per-technique statistics from identity labels and eval-hash history. + + Args: + technique_identifiers (Sequence[str]): Stable technique identifiers to + aggregate. Returned dict is keyed by these identifiers. + label_name (str): Result-label key containing the technique identifier. + technique_eval_hashes_by_identifier (Mapping[str, str] | None): + Optional mapping from requested selector identifiers to the full + ``AttackTechnique`` eval hashes persisted by normal scenarios. + Matching labeled and eval-hash rows are merged by result ID so a + row visible through both paths is counted once. + scenario_result_id (str | None): Restrict to a single scenario run. + Defaults to ``None`` (aggregate across all runs). + targeted_harm_categories (Sequence[str] | None): Restrict to results + whose attack targeted these harm categories. Defaults to ``None``. + memory (MemoryInterface | None): Memory backend to query. Defaults to + ``CentralMemory.get_memory_instance()``. - s, f, u, e = counts.get(eval_hash, (0, 0, 0, 0)) - if result.outcome == AttackOutcome.SUCCESS: - counts[eval_hash] = (s + 1, f, u, e) - elif result.outcome == AttackOutcome.FAILURE: - counts[eval_hash] = (s, f + 1, u, e) - elif result.outcome == AttackOutcome.ERROR: - counts[eval_hash] = (s, f, u, e + 1) + Returns: + dict[str, AttackStats]: Stats per requested technique identifier. + Identifiers with no historical results are omitted. + """ + if not technique_identifiers: + return {} + + if memory is None: + memory = CentralMemory.get_memory_instance() + labeled_results = memory.get_attack_results( + labels={label_name: list(technique_identifiers)}, + scenario_result_id=scenario_result_id, + targeted_harm_categories=targeted_harm_categories, + ) + eval_results = ( + memory.get_attack_results( + atomic_attack_eval_hashes=sorted(set(technique_eval_hashes_by_identifier.values())), + scenario_result_id=scenario_result_id, + targeted_harm_categories=targeted_harm_categories, + ) + if technique_eval_hashes_by_identifier + else [] + ) + + requested = set(technique_identifiers) + identifiers_by_eval_hash: dict[str, list[str]] = {} + for technique_identifier, eval_hash in (technique_eval_hashes_by_identifier or {}).items(): + if technique_identifier in requested: + identifiers_by_eval_hash.setdefault(eval_hash, []).append(technique_identifier) + + unique_results = {result.attack_result_id: result for result in [*labeled_results, *eval_results]} + grouped_outcomes: list[tuple[str, AttackOutcome]] = [] + for result in unique_results.values(): + labeled_identifier = result.labels.get(label_name) + if labeled_identifier in requested: + matching_identifiers = [labeled_identifier] else: - counts[eval_hash] = (s, f, u + 1, e) + result_identifier = result.atomic_attack_identifier + result_eval_hash = result_identifier.eval_hash if result_identifier is not None else None + matching_identifiers = identifiers_by_eval_hash.get(result_eval_hash or "", []) + if not matching_identifiers: + continue - return { - eval_hash: _compute_stats(successes=s, failures=f, undetermined=u, errors=e) - for eval_hash, (s, f, u, e) in counts.items() - } + grouped_outcomes.extend((technique_identifier, result.outcome) for technique_identifier in matching_identifiers) + + return _compute_grouped_outcome_stats(grouped_outcomes) diff --git a/pyrit/backend/services/scenario_configuration_resolver.py b/pyrit/backend/services/scenario_configuration_resolver.py index a0310bf595..4b967d1fa3 100644 --- a/pyrit/backend/services/scenario_configuration_resolver.py +++ b/pyrit/backend/services/scenario_configuration_resolver.py @@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Any from pyrit.registry import ConverterRegistry, ScenarioRegistry, TargetRegistry +from pyrit.scenario.core.dataset_configuration import CompoundDatasetAttackConfiguration if TYPE_CHECKING: from pyrit.converter import Converter @@ -125,7 +126,30 @@ def resolve_configuration( if dataset_names or max_dataset_size is not None or filters: default_config = introspection_instance._default_dataset_config - if dataset_names: + + if isinstance(default_config, CompoundDatasetAttackConfiguration): + names_changed = dataset_names is not None and dataset_names != default_config.dataset_names + if names_changed: + try: + resolved["dataset_config"] = default_config.with_dataset_names( + dataset_names=dataset_names, + max_dataset_size=max_dataset_size, + filters=filters or None, + ) + except TypeError as exc: + raise ValueError( + f"Scenario '{scenario_name}' does not support overriding datasets through " + f"its {type(default_config).__name__} configuration: {exc}" + ) from exc + else: + if max_dataset_size is not None: + default_config.update_child_max_dataset_size(max_dataset_size=max_dataset_size) + if filters: + default_config.update_filters(filters=filters) + resolved["dataset_config"] = default_config + elif dataset_names: + # Construct a fresh instance of the scenario's own dataset-config + # class so subclass-specific behavior is preserved. default_config_class = type(default_config) try: resolved["dataset_config"] = default_config_class( @@ -139,6 +163,10 @@ def resolve_configuration( f"its {default_config_class.__name__} configuration: {exc}" ) from exc else: + # Reuse the scenario's default dataset config (preserves subtype + + # the scenario's own default dataset names) and override only the + # sample cap and/or filters. Safe because the introspection instance + # is throwaway. if max_dataset_size is not None: default_config.max_dataset_size = max_dataset_size if filters: diff --git a/pyrit/backend/services/scenario_service.py b/pyrit/backend/services/scenario_service.py index 4d52bce35a..eb0b7f9e63 100644 --- a/pyrit/backend/services/scenario_service.py +++ b/pyrit/backend/services/scenario_service.py @@ -14,8 +14,9 @@ from pyrit.backend.models.common import PaginationInfo from pyrit.backend.models.scenarios import ListRegisteredScenariosResponse from pyrit.backend.services.scenario_configuration_resolver import ScenarioConfigurationResolver -from pyrit.models.catalog.scenario import ( +from pyrit.models.catalog import ( RegisteredScenario, + ScenarioDefaultRunSizeEstimate, ScenarioRunSizeEstimate, ScenarioRunSizeEstimateRequest, ) @@ -23,7 +24,9 @@ logger = logging.getLogger(__name__) _ESTIMATE_CACHE_SIZE = 128 -_ESTIMATE_CONCURRENCY = 1 +_ESTIMATE_CONCURRENCY = 4 +_CONFIGURED_ESTIMATE_CONCURRENCY = 4 +_DEFAULT_ESTIMATE_TIMEOUT_SECONDS = 3.0 _ESTIMATE_INFLIGHT_SIZE = 256 _UNAVAILABLE_CACHE_TTL_SECONDS = 30.0 _EstimateCacheKey = tuple[str, int] @@ -62,6 +65,7 @@ def _metadata_to_registered_scenario( all_techniques=list(metadata.all_techniques), technique_summaries=list(metadata.technique_summaries), default_datasets=list(metadata.default_datasets), + dataset_size_limit=metadata.dataset_size_limit, supported_parameters=list(metadata.supported_parameters), baseline_policy=metadata.baseline_policy, include_baseline_by_default=metadata.include_baseline_by_default, @@ -79,6 +83,7 @@ def __init__(self) -> None: self._estimate_tasks: OrderedDict[_EstimateCacheKey, _EstimateTask] = OrderedDict() self._estimate_task_lock = asyncio.Lock() self._estimate_semaphore = asyncio.Semaphore(_ESTIMATE_CONCURRENCY) + self._configured_estimate_semaphore = asyncio.Semaphore(_CONFIGURED_ESTIMATE_CONCURRENCY) async def list_scenarios_async( self, @@ -162,11 +167,7 @@ async def estimate_scenario_run_size_async( if metadata is None: return None - semaphore = getattr(self, "_estimate_semaphore", None) - if semaphore is None: - semaphore = asyncio.Semaphore(_ESTIMATE_CONCURRENCY) - self._estimate_semaphore = semaphore - async with semaphore: + async with self._configured_estimate_semaphore: return await self._estimate_configured_run_size_async( scenario_name=scenario_name, request=request, @@ -175,30 +176,19 @@ async def estimate_scenario_run_size_async( async def _get_default_run_size_estimate_async(self, *, metadata: ScenarioMetadata) -> ScenarioRunSizeEstimate: """Return a cached, cancellation-safe scenario-owned estimate.""" cache_key = (metadata.registry_name, metadata.scenario_version) - cache = getattr(self, "_estimate_cache", None) - if cache is None: - cache = OrderedDict() - self._estimate_cache = cache while True: cached = self._read_estimate_cache(cache_key=cache_key) if cached is not None: return cached - task_lock = getattr(self, "_estimate_task_lock", None) - if task_lock is None: - task_lock = asyncio.Lock() - self._estimate_task_lock = task_lock wait_for_capacity: _EstimateTask | None = None task: _EstimateTask | None = None - async with task_lock: + async with self._estimate_task_lock: cached = self._read_estimate_cache(cache_key=cache_key) if cached is not None: return cached - tasks = getattr(self, "_estimate_tasks", None) - if tasks is None: - tasks = OrderedDict() - self._estimate_tasks = tasks + tasks = self._estimate_tasks for completed_key in [key for key, candidate in tasks.items() if candidate.done()]: del tasks[completed_key] task = tasks.get(cache_key) @@ -251,19 +241,26 @@ async def _compute_default_run_size_estimate_async( Returns: ScenarioRunSizeEstimate: Scenario-owned estimate. """ - semaphore = getattr(self, "_estimate_semaphore", None) - if semaphore is None: - semaphore = asyncio.Semaphore(_ESTIMATE_CONCURRENCY) - self._estimate_semaphore = semaphore - async with semaphore: - try: - scenario = await asyncio.to_thread(self._registry.create_instance, scenario_name) - estimate = await scenario.get_default_run_size_estimate_async() - except Exception as exc: - logger.warning("Default-run estimate failed for scenario '%s': %s", scenario_name, exc) - estimate = ScenarioRunSizeEstimate.unavailable( - note=f"The scenario could not resolve its default inputs for estimation ({type(exc).__name__})." + try: + async with self._estimate_semaphore: + estimate = await asyncio.wait_for( + self._run_default_estimate_async(scenario_name=scenario_name), + timeout=_DEFAULT_ESTIMATE_TIMEOUT_SECONDS, ) + except TimeoutError: + logger.warning( + "Default-run estimate timed out for scenario '%s' after %.1f seconds", + scenario_name, + _DEFAULT_ESTIMATE_TIMEOUT_SECONDS, + ) + estimate = ScenarioDefaultRunSizeEstimate.unavailable( + note="The default estimate timed out; open the scenario to calculate the configured run size." + ) + except Exception as exc: + logger.warning("Default-run estimate failed for scenario '%s': %s", scenario_name, exc) + estimate = ScenarioDefaultRunSizeEstimate.unavailable( + note=f"The scenario could not resolve its default inputs for estimation ({type(exc).__name__})." + ) expires_at = monotonic() + _UNAVAILABLE_CACHE_TTL_SECONDS if estimate.estimated_attack_count is None else None cache = self._estimate_cache @@ -273,6 +270,16 @@ async def _compute_default_run_size_estimate_async( cache.popitem(last=False) return estimate + async def _run_default_estimate_async(self, *, scenario_name: str) -> ScenarioDefaultRunSizeEstimate: + """ + Run one default estimate. + + Returns: + ScenarioDefaultRunSizeEstimate: The authoritative scenario estimate. + """ + scenario = await asyncio.to_thread(self._registry.create_instance, scenario_name) + return await scenario.get_default_run_size_estimate_async() + def _clear_estimate_task(self, *, task: _EstimateTask, cache_key: _EstimateCacheKey) -> None: """Remove a completed single-flight task without disturbing a replacement.""" tasks = self._estimate_tasks diff --git a/pyrit/models/__init__.py b/pyrit/models/__init__.py index 5b8d2b8bdc..9d867dcdde 100644 --- a/pyrit/models/__init__.py +++ b/pyrit/models/__init__.py @@ -24,11 +24,14 @@ if TYPE_CHECKING: from pyrit.models.additional_initializer import AdditionalInitializer from pyrit.models.catalog import ( + ScenarioAdaptiveRunSizeDetails, ScenarioDatasetSizeCap, + ScenarioDatasetSizeLimit, ScenarioDatasetSummary, ScenarioRunListItem, ScenarioRunSizeComponent, ScenarioRunSizeEstimate, + ScenarioRunSizeEstimateCondition, ScenarioRunSizeEstimateRequest, ScenarioTechniqueSummary, ) @@ -252,11 +255,14 @@ "ScorerEvaluationIdentifier": "pyrit.models.identifiers", "ScorerIdentifier": "pyrit.models.identifiers", "ScenarioIdentifier": "pyrit.models.identifiers", + "ScenarioAdaptiveRunSizeDetails": "pyrit.models.catalog", "ScenarioDatasetSizeCap": "pyrit.models.catalog", + "ScenarioDatasetSizeLimit": "pyrit.models.catalog", "ScenarioDatasetSummary": "pyrit.models.catalog", "ScenarioRunListItem": "pyrit.models.catalog", "ScenarioRunSizeComponent": "pyrit.models.catalog", "ScenarioRunSizeEstimate": "pyrit.models.catalog", + "ScenarioRunSizeEstimateCondition": "pyrit.models.catalog", "ScenarioRunSizeEstimateRequest": "pyrit.models.catalog", "ScenarioTechniqueSummary": "pyrit.models.catalog", "ScenarioResult": "pyrit.models.results.scenario_result", diff --git a/pyrit/models/catalog/__init__.py b/pyrit/models/catalog/__init__.py index 18b3f6b713..a6b4794e00 100644 --- a/pyrit/models/catalog/__init__.py +++ b/pyrit/models/catalog/__init__.py @@ -25,13 +25,21 @@ AttackRetrySummary, RegisteredScenario, RunScenarioRequest, + ScenarioAdaptiveRunSizeDetails, ScenarioDatasetSizeCap, + ScenarioDatasetSizeLimit, ScenarioDatasetSummary, + ScenarioDefaultRunSizeEstimate, + ScenarioOverloadSummary, ScenarioRunListItem, ScenarioRunSizeComponent, ScenarioRunSizeEstimate, + ScenarioRunSizeEstimateCondition, ScenarioRunSizeEstimateRequest, + ScenarioRunSizeEstimateStatus, + ScenarioRunSizeFactor, ScenarioRunSummary, + ScenarioTargetSummary, ScenarioTechniqueSummary, ) from pyrit.models.catalog.target import TargetInstance @@ -42,14 +50,22 @@ "RegisteredInitializer": "pyrit.models.catalog.initializer", "RegisteredScenario": "pyrit.models.catalog.scenario", "RunScenarioRequest": "pyrit.models.catalog.scenario", + "ScenarioAdaptiveRunSizeDetails": "pyrit.models.catalog.scenario", "ScenarioDatasetSizeCap": "pyrit.models.catalog.scenario", + "ScenarioDatasetSizeLimit": "pyrit.models.catalog.scenario", "ScenarioDatasetSummary": "pyrit.models.catalog.scenario", + "ScenarioDefaultRunSizeEstimate": "pyrit.models.catalog.scenario", + "ScenarioOverloadSummary": "pyrit.models.catalog.scenario", "ScenarioRunListItem": "pyrit.models.catalog.scenario", "ScenarioRunSizeComponent": "pyrit.models.catalog.scenario", "ScenarioRunSizeEstimate": "pyrit.models.catalog.scenario", + "ScenarioRunSizeEstimateCondition": "pyrit.models.catalog.scenario", "ScenarioRunSizeEstimateRequest": "pyrit.models.catalog.scenario", + "ScenarioRunSizeEstimateStatus": "pyrit.models.catalog.scenario", + "ScenarioRunSizeFactor": "pyrit.models.catalog.scenario", "ScenarioRunSummary": "pyrit.models.catalog.scenario", "ScenarioTechniqueSummary": "pyrit.models.catalog.scenario", + "ScenarioTargetSummary": "pyrit.models.catalog.scenario", "TargetInstance": "pyrit.models.catalog.target", } diff --git a/pyrit/models/catalog/scenario.py b/pyrit/models/catalog/scenario.py index d46af4068c..221e26f769 100644 --- a/pyrit/models/catalog/scenario.py +++ b/pyrit/models/catalog/scenario.py @@ -14,6 +14,8 @@ """ from datetime import datetime +from enum import Enum +from math import prod from typing import Any, Literal from pydantic import AliasChoices, BaseModel, Field, field_validator, model_validator @@ -57,14 +59,116 @@ def _validate_dataset_filter_mapping( return value +class ScenarioRunSizeEstimateStatus(str, Enum): + """Confidence level for a catalog default-run size estimate.""" + + Exact = "exact" + Conditional = "conditional" + Unavailable = "unavailable" + + +class ScenarioRunSizeEstimateCondition(str, Enum): + """Reason an estimate remains conditional until launch.""" + + TargetCapabilities = "target_capabilities" + LaunchConfiguration = "launch_configuration" + + +class ScenarioRunSizeFactor(BaseModel): + """One labeled multiplicative factor in a run-size component.""" + + label: str = Field(..., min_length=1) + count: int = Field(..., ge=0) + + class ScenarioRunSizeComponent(BaseModel): """One additive component of a default-run size estimate.""" label: str = Field(..., min_length=1) count: int = Field(..., ge=0) + factors: list[ScenarioRunSizeFactor] = Field(default_factory=list) is_baseline: bool = False + condition: ScenarioRunSizeEstimateCondition | None = None note: str | None = None + @model_validator(mode="after") + def validate_factor_product(self) -> "ScenarioRunSizeComponent": + """ + Require known component totals to equal their ordered factor product. + + Returns: + ScenarioRunSizeComponent: The validated component. + + Raises: + ValueError: If a component with factors has an inconsistent count. + """ + if self.factors: + factor_product = prod(factor.count for factor in self.factors) + if self.count != factor_product: + raise ValueError( + f"Component '{self.label}' count ({self.count}) must equal its factor product ({factor_product})" + ) + return self + + +class ScenarioAdaptiveRunSizeDetails(BaseModel): + """Structured work bounds for an adaptive scenario estimate.""" + + objective_count: int = Field(..., ge=0) + selected_candidate_technique_count: int = Field(..., ge=1) + candidate_technique_count: int = Field(..., ge=1) + max_attempts_per_objective: int = Field(..., ge=1) + techniques_per_objective_upper_bound: int = Field(..., ge=1) + technique_attempt_count_upper_bound: int = Field(..., ge=0) + stop_on_first_success: Literal[True] = True + compatibility_may_reduce_attempts: Literal[True] = True + + @model_validator(mode="before") + @classmethod + def default_selected_candidate_count(cls, data: Any) -> Any: + """ + Preserve version-1 payload compatibility when the selected count is absent. + + Returns: + Any: Input data with the selected count defaulted to the compatible count. + """ + if ( + isinstance(data, dict) + and "selected_candidate_technique_count" not in data + and "candidate_technique_count" in data + ): + return { + **data, + "selected_candidate_technique_count": data["candidate_technique_count"], + } + return data + + @model_validator(mode="after") + def validate_attempt_bounds(self) -> "ScenarioAdaptiveRunSizeDetails": + """ + Ensure the serialized adaptive attempt bounds match the configured pool. + + Returns: + ScenarioAdaptiveRunSizeDetails: The validated details. + + Raises: + ValueError: If either derived upper bound is inconsistent. + """ + if self.candidate_technique_count > self.selected_candidate_technique_count: + raise ValueError("candidate_technique_count cannot exceed selected_candidate_technique_count") + expected_per_objective = min(self.candidate_technique_count, self.max_attempts_per_objective) + if self.techniques_per_objective_upper_bound != expected_per_objective: + raise ValueError( + "techniques_per_objective_upper_bound must equal " + "min(candidate_technique_count, max_attempts_per_objective)" + ) + expected_total = self.objective_count * expected_per_objective + if self.technique_attempt_count_upper_bound != expected_total: + raise ValueError( + "technique_attempt_count_upper_bound must equal objective_count * techniques_per_objective_upper_bound" + ) + return self + class ScenarioDatasetSizeCap(BaseModel): """One configured cap affecting a dataset or compound population.""" @@ -98,7 +202,31 @@ class ScenarioTechniqueSummary(BaseModel): tags: list[str] = Field(default_factory=list) -class ScenarioRunSizeEstimate(BaseModel): +class ScenarioDatasetSizeLimit(BaseModel): + """Structured default and override semantics for a scenario's dataset-size limit.""" + + default_scope: Literal["none", "per_dataset", "combined", "heterogeneous"] = "none" + default_count: int | None = Field(default=None, ge=1) + override_scope: Literal["per_dataset", "combined", "unsupported"] = "per_dataset" + + @model_validator(mode="after") + def validate_default_count(self) -> "ScenarioDatasetSizeLimit": + """ + Require a count exactly when the default has one representable scope. + + Returns: + ScenarioDatasetSizeLimit: The validated limit metadata. + + Raises: + ValueError: If the count does not match the declared default scope. + """ + has_representable_default = self.default_scope in {"per_dataset", "combined"} + if has_representable_default != (self.default_count is not None): + raise ValueError("default_count must be set exactly for per_dataset or combined defaults") + return self + + +class ScenarioDefaultRunSizeEstimate(BaseModel): """ Structured estimate of default planned scenario execution units. @@ -106,52 +234,129 @@ class ScenarioRunSizeEstimate(BaseModel): logical-seed-group pair. Retries and internal attack turns are excluded. """ - estimated_attack_count: int | None = Field(default=None, ge=0) + version: Literal[1] = 1 + status: ScenarioRunSizeEstimateStatus + total_attack_count: int | None = Field( + default=None, + ge=0, + validation_alias=AliasChoices("total_attack_count", "total", "estimated_attack_count"), + ) minimum_attack_count: int | None = Field(default=None, ge=0) maximum_attack_count: int | None = Field(default=None, ge=0) + condition: ScenarioRunSizeEstimateCondition | None = None components: list[ScenarioRunSizeComponent] = Field(default_factory=list) datasets: list[ScenarioDatasetSummary] = Field(default_factory=list) + adaptive_details: ScenarioAdaptiveRunSizeDetails | None = None effective_parameters: dict[str, bool | int | float | str | list[str]] = Field( default_factory=dict, description="Scenario parameter values used by this estimate, including implicit runtime defaults.", ) - note: str | None = None + note: str | None = Field(default=None, validation_alias=AliasChoices("note", "caveat")) + retries_included: Literal[False] = False + + @property + def total(self) -> int | None: + """The legacy Python attribute for total_attack_count.""" + return self.total_attack_count + + @property + def estimated_attack_count(self) -> int | None: + """The legacy Python attribute for total_attack_count.""" + return self.total_attack_count + + @property + def caveat(self) -> str | None: + """The legacy Python attribute for note.""" + return self.note + + @model_validator(mode="before") + @classmethod + def normalize_legacy_estimate(cls, data: Any) -> Any: + """ + Default an omitted status from the available total, then explain + component-less legacy exact totals in the canonical shape. + + Callers that predate the ``status`` field (e.g. ``estimated_attack_count``-only + constructors) are treated as ``Exact`` when they supply a total and + ``Conditional`` otherwise. + + Returns: + Any: The normalized input. + """ + if not isinstance(data, dict): + return data + normalized = dict(data) + if "status" not in normalized: + has_total = any( + normalized.get(key) is not None for key in ("total_attack_count", "total", "estimated_attack_count") + ) + normalized["status"] = ( + ScenarioRunSizeEstimateStatus.Exact if has_total else ScenarioRunSizeEstimateStatus.Conditional + ) + if "total" in normalized and "components" not in normalized: + status = normalized.get("status") + if status == ScenarioRunSizeEstimateStatus.Exact or status == "exact": + normalized["components"] = [ + { + "label": "Legacy total", + "count": normalized["total"], + "note": "Normalized from a legacy component-less estimate.", + } + ] + return normalized @model_validator(mode="after") - def validate_estimated_attack_count(self) -> "ScenarioRunSizeEstimate": + def validate_total(self) -> "ScenarioDefaultRunSizeEstimate": """ - Ensure available estimates expose a complete additive total. + Ensure exact estimates expose and explain their complete total. Returns: - ScenarioRunSizeEstimate: The validated estimate. + ScenarioDefaultRunSizeEstimate: The validated estimate. Raises: - ValueError: If an available estimate misstates its total. + ValueError: If an exact estimate omits or misstates its total. """ if ( self.minimum_attack_count is not None and self.maximum_attack_count is not None and self.minimum_attack_count > self.maximum_attack_count ): - raise ValueError("Minimum attack count cannot exceed maximum attack count") + raise ValueError("minimum_attack_count must be less than or equal to maximum_attack_count") - if self.estimated_attack_count is not None: - component_total = sum(component.count for component in self.components) - if component_total != self.estimated_attack_count: - raise ValueError( - f"Default-run estimate components total {component_total}, not {self.estimated_attack_count}" - ) + if self.status is not ScenarioRunSizeEstimateStatus.Exact: + return self + + if self.total_attack_count is None: + raise ValueError("Exact default-run estimates require total_attack_count") + for field_name, bound in ( + ("minimum_attack_count", self.minimum_attack_count), + ("maximum_attack_count", self.maximum_attack_count), + ): + if bound is not None and bound != self.total_attack_count: + raise ValueError(f"Exact default-run estimates require {field_name} to equal total_attack_count") + + component_total = sum(component.count for component in self.components) + if component_total != self.total_attack_count: + raise ValueError( + f"Exact default-run estimate components total {component_total}, not {self.total_attack_count}" + ) return self @classmethod - def unavailable(cls, *, note: str = "Default-run size estimate is unavailable.") -> "ScenarioRunSizeEstimate": + def unavailable( + cls, *, note: str = "Default-run size estimate is unavailable." + ) -> "ScenarioDefaultRunSizeEstimate": """ Build an unavailable estimate without presenting a guessed total. Returns: - ScenarioRunSizeEstimate: An unavailable estimate. + ScenarioDefaultRunSizeEstimate: An unavailable estimate. """ - return cls(note=note) + return cls(status=ScenarioRunSizeEstimateStatus.Unavailable, note=note) + + +# Backward-compatible catalog name from the initial run-size DTO. +ScenarioRunSizeEstimate = ScenarioDefaultRunSizeEstimate class RegisteredScenario(BaseModel): @@ -186,6 +391,10 @@ 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_size_limit: ScenarioDatasetSizeLimit = Field( + default_factory=ScenarioDatasetSizeLimit, + description="Structured scenario-default and explicit-override dataset-size limit semantics", + ) baseline_policy: Literal["enabled", "disabled", "forbidden"] = Field( "enabled", description="Whether baseline execution is enabled, disabled, or forbidden" ) diff --git a/pyrit/registry/components/scenario_registry.py b/pyrit/registry/components/scenario_registry.py index cbd84f5206..cca2a1c576 100644 --- a/pyrit/registry/components/scenario_registry.py +++ b/pyrit/registry/components/scenario_registry.py @@ -18,7 +18,12 @@ 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 ( + ScenarioDatasetSizeLimit, + 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 +71,9 @@ class ScenarioMetadata(RegistryMetadata): # Default dataset names used by this scenario. default_datasets: tuple[str, ...] = field(kw_only=True) + # Structured default and override semantics for the dataset-size control. + dataset_size_limit: ScenarioDatasetSizeLimit = field(kw_only=True, default_factory=ScenarioDatasetSizeLimit) + # Scenario-declared custom parameters. supported_parameters: tuple[Parameter, ...] = field(kw_only=True, default=()) @@ -193,7 +201,12 @@ def _build_metadata(self, name: str, cls: type[Scenario]) -> ScenarioMetadata: ) for aggregate in technique_class.get_aggregate_techniques() ) - default_datasets = tuple(instance._default_dataset_config.dataset_names) + default_dataset_config = instance._default_dataset_config + default_datasets = tuple(default_dataset_config.dataset_names) + dataset_size_limit = self._build_dataset_size_limit( + default_dataset_config=default_dataset_config, + override_scope=instance.get_dataset_size_limit_override_scope(), + ) return ScenarioMetadata( class_name=cls.__name__, @@ -209,11 +222,46 @@ 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_size_limit=dataset_size_limit, supported_parameters=supported_parameters, baseline_policy=instance.BASELINE_ATTACK_POLICY.value, include_baseline_by_default=instance.BASELINE_ATTACK_POLICY.value == "enabled", ) + @staticmethod + def _build_dataset_size_limit( + *, + default_dataset_config: Any, + override_scope: Literal["per_dataset", "combined", "unsupported"], + ) -> ScenarioDatasetSizeLimit: + """ + Normalize a scenario dataset configuration into form-level limit semantics. + + Returns: + ScenarioDatasetSizeLimit: The structured default and override scopes. + """ + dataset_names = tuple(default_dataset_config.dataset_names) + caps_by_dataset = default_dataset_config.size_caps_by_dataset() + configured_caps = [caps for caps in caps_by_dataset.values() if caps] + if not configured_caps: + return ScenarioDatasetSizeLimit(default_scope="none", override_scope=override_scope) + + expected_source_count = max(1, len(set(dataset_names))) + if all(len(caps) == 1 for caps in configured_caps) and len(configured_caps) == expected_source_count: + cap_signatures = {(caps[0][1], caps[0][2]) for caps in configured_caps} + if len(cap_signatures) == 1: + count, configured_on = next(iter(cap_signatures)) + default_scope: Literal["per_dataset", "combined"] = ( + "per_dataset" if configured_on == "dataset" else "combined" + ) + return ScenarioDatasetSizeLimit( + default_scope=default_scope, + default_count=count, + override_scope=override_scope, + ) + + return ScenarioDatasetSizeLimit(default_scope="heterogeneous", override_scope=override_scope) + async def create_and_estimate_async( self, *, diff --git a/pyrit/scenario/core/dataset_configuration.py b/pyrit/scenario/core/dataset_configuration.py index 4f43d5fdc6..4e8261b001 100644 --- a/pyrit/scenario/core/dataset_configuration.py +++ b/pyrit/scenario/core/dataset_configuration.py @@ -338,9 +338,10 @@ def __init__( self._dataset_names = list(dataset_names) if dataset_names is not None else None self.max_dataset_size = max_dataset_size self._filters: dict[str, list[str]] = dict(filters or {}) + self._custom_validators = list(validators) if validators else [] self._validators: list[Callable[[ResolvedDataset], None]] = [ *self._default_validators(), - *(list(validators) if validators else []), + *self._custom_validators, ] self._auto_fetch = auto_fetch @@ -742,6 +743,28 @@ async def get_attack_groups_by_dataset_async( raise DatasetConstraintError(f"Resolved attack-group dataset is empty (datasets: {names}).") return result + async def resolve_attack_groups_for_estimate_async( + self, + ) -> tuple[dict[str, list[AttackSeedGroup]], dict[str, list[AttackSeedGroup]]]: + """ + Resolve full and sampled attack groups with one dataset fetch. + + Returns: + tuple: Full groups and effectively selected groups, both keyed by dataset. + + Raises: + DatasetConstraintError: If the resolved or sampled attack-group population is empty. + """ + groups_by_dataset, resolved = await self._build_groups_by_dataset_async() + self.validate(resolved) + selected = { + name: groups for name, groups in self._sample_groups_by_dataset(groups_by_dataset).items() if groups + } + if not groups_by_dataset or not selected: + names = ", ".join(self._dataset_names) if self._dataset_names else "" + raise DatasetConstraintError(f"Resolved attack-group dataset is empty (datasets: {names}).") + return groups_by_dataset, selected + def _sample_groups_by_dataset( self, groups_by_dataset: dict[str, list[AttackSeedGroup]] ) -> dict[str, list[AttackSeedGroup]]: @@ -849,6 +872,75 @@ def per_dataset( ] ) + def with_dataset_names( + self, + *, + dataset_names: Sequence[str], + max_dataset_size: int | None = None, + filters: dict[str, list[str]] | None = None, + ) -> CompoundDatasetAttackConfiguration: + """ + Rebuild a homogeneous per-dataset compound for an explicit name selection. + + This preserves the scenario's per-dataset cap, auto-fetch policy, and shared + filters when every child is a plain single-dataset attack configuration. + Heterogeneous compounds must provide their own scenario-specific override + path rather than silently losing child shaping behavior. + + Args: + dataset_names (Sequence[str]): Selected dataset names in request order. + max_dataset_size (int | None): Optional replacement per-dataset cap. + filters (dict[str, list[str]] | None): Filters merged over the shared defaults. + + Returns: + CompoundDatasetAttackConfiguration: A fresh compound for the selected datasets. + + Raises: + TypeError: If the compound has heterogeneous or shaped child configurations. + ValueError: If ``dataset_names`` is empty or contains duplicates. + """ + if len(set(dataset_names)) != len(dataset_names): + raise ValueError("dataset-name overrides cannot contain duplicates") + if any( + type(child) is not DatasetAttackConfiguration or len(child.dataset_names) != 1 + for child in self._configurations + ): + raise TypeError( + "dataset-name overrides require homogeneous single-dataset DatasetAttackConfiguration children" + ) + + child_caps = {child.max_dataset_size for child in self._configurations} + child_auto_fetch = {child._auto_fetch for child in self._configurations} + child_filters = { + tuple(sorted((key, tuple(values)) for key, values in child.filters.items())) + for child in self._configurations + } + template_child = self._configurations[0] + child_validators_match = all( + child._custom_validators == template_child._custom_validators for child in self._configurations[1:] + ) + if len(child_caps) != 1 or len(child_auto_fetch) != 1 or len(child_filters) != 1 or not child_validators_match: + raise TypeError( + "dataset-name overrides require children with shared caps, filters, validators, and auto-fetch policy" + ) + + inherited_filters = {key: list(values) for key, values in next(iter(child_filters))} + inherited_filters.update(filters or {}) + per_dataset_cap = max_dataset_size if max_dataset_size is not None else next(iter(child_caps)) + rebuilt = type(self).per_dataset( + dataset_names=dataset_names, + max_dataset_size=per_dataset_cap, + auto_fetch=next(iter(child_auto_fetch)), + filters=inherited_filters or None, + ) + for child in rebuilt._configurations: + child._custom_validators = list(template_child._custom_validators) + child._validators = list(template_child._validators) + rebuilt.max_dataset_size = self.max_dataset_size + rebuilt._custom_validators = list(self._custom_validators) + rebuilt._validators = list(self._validators) + return rebuilt + @property def dataset_names(self) -> list[str]: """ @@ -911,6 +1003,21 @@ def update_filters(self, *, filters: dict[str, list[str]]) -> None: for child in self._configurations: child.update_filters(filters=filters) + def update_child_max_dataset_size(self, *, max_dataset_size: int) -> None: + """ + Apply the same independent sampling cap to every child configuration. + + Args: + max_dataset_size (int): Positive per-child logical-group cap. + + Raises: + ValueError: If ``max_dataset_size`` is less than one. + """ + if max_dataset_size < 1: + raise ValueError("'max_dataset_size' must be a positive integer (>= 1).") + for child in self._configurations: + child.max_dataset_size = max_dataset_size + async def get_attack_seed_groups_async(self, *, apply_sampling: bool = True) -> list[AttackSeedGroup]: """ Concatenate every child's flat result, then validate and apply the global cap. @@ -961,6 +1068,27 @@ async def get_attack_groups_by_dataset_async( self.validate(self._resolved_from_groups([group for groups in merged.values() for group in groups])) return self._sample_groups_by_dataset(merged) if apply_sampling else merged + async def resolve_attack_groups_for_estimate_async( + self, + ) -> tuple[dict[str, list[AttackSeedGroup]], dict[str, list[AttackSeedGroup]]]: + """ + Resolve every child's full and sampled populations with one fetch per child. + + Returns: + tuple: Full groups and effectively selected groups, both keyed by dataset. + """ + full_merged: dict[str, list[AttackSeedGroup]] = {} + selected_merged: dict[str, list[AttackSeedGroup]] = {} + for child in self._configurations: + full_groups, selected_groups = await child.resolve_attack_groups_for_estimate_async() + for name, groups in full_groups.items(): + full_merged.setdefault(name, []).extend(groups) + for name, groups in selected_groups.items(): + selected_merged.setdefault(name, []).extend(groups) + self.validate(self._resolved_from_groups([group for groups in full_merged.values() for group in groups])) + selected = {name: groups for name, groups in self._sample_groups_by_dataset(selected_merged).items() if groups} + return full_merged, selected + def _resolved_from_groups(self, groups: list[AttackSeedGroup]) -> ResolvedDataset: """ Build a ResolvedDataset over the combined groups for compound-level validation. diff --git a/pyrit/scenario/core/scenario.py b/pyrit/scenario/core/scenario.py index 5e36c15d46..70be3461fd 100644 --- a/pyrit/scenario/core/scenario.py +++ b/pyrit/scenario/core/scenario.py @@ -795,10 +795,13 @@ async def _resolve_dataset_groups_for_estimate_async( configured_dataset = self._dataset_config with read_only_dataset_resolution(): self._dataset_config = configured_dataset - full_groups = await self._resolve_seed_groups_by_dataset_async(apply_sampling=False) + if type(self)._resolve_seed_groups_by_dataset_async is Scenario._resolve_seed_groups_by_dataset_async: + full_groups, selected_groups = await configured_dataset.resolve_attack_groups_for_estimate_async() + else: + full_groups = await self._resolve_seed_groups_by_dataset_async(apply_sampling=False) + self._dataset_config = configured_dataset + selected_groups = await self._resolve_seed_groups_by_dataset_async(apply_sampling=True) self._estimate_full_groups_by_dataset = full_groups - self._dataset_config = configured_dataset - selected_groups = await self._resolve_seed_groups_by_dataset_async(apply_sampling=True) configured_caps = self._dataset_config.size_caps_by_dataset() datasets: list[ScenarioDatasetSummary] = [] diff --git a/pyrit/scenario/scenarios/adaptive/__init__.py b/pyrit/scenario/scenarios/adaptive/__init__.py index b71c7fb445..a74196af63 100644 --- a/pyrit/scenario/scenarios/adaptive/__init__.py +++ b/pyrit/scenario/scenarios/adaptive/__init__.py @@ -12,6 +12,8 @@ from pyrit.scenario.scenarios.adaptive.adaptive_scenario import AdaptiveScenario from pyrit.scenario.scenarios.adaptive.dispatcher import ( ADAPTIVE_ATTEMPT_LABEL, + ADAPTIVE_TECHNIQUE_ID_LABEL, + ADAPTIVE_TECHNIQUE_NAME_LABEL, AdaptiveTechniqueDispatcher, TechniqueBundle, ) @@ -20,11 +22,15 @@ SelectorScope, TechniqueSelector, ) + from pyrit.scenario.scenarios.adaptive.technique_identity import AdaptiveTechniqueIdentifier from pyrit.scenario.scenarios.adaptive.text_adaptive import TextAdaptive _LAZY_EXPORTS: dict[str, str | tuple[str, str | None]] = { "ADAPTIVE_ATTEMPT_LABEL": "pyrit.scenario.scenarios.adaptive.dispatcher", + "ADAPTIVE_TECHNIQUE_ID_LABEL": "pyrit.scenario.scenarios.adaptive.dispatcher", + "ADAPTIVE_TECHNIQUE_NAME_LABEL": "pyrit.scenario.scenarios.adaptive.dispatcher", "AdaptiveScenario": "pyrit.scenario.scenarios.adaptive.adaptive_scenario", + "AdaptiveTechniqueIdentifier": "pyrit.scenario.scenarios.adaptive.technique_identity", "AdaptiveTechniqueDispatcher": "pyrit.scenario.scenarios.adaptive.dispatcher", "EpsilonGreedyTechniqueSelector": "pyrit.scenario.scenarios.adaptive.selectors", "SelectorScope": "pyrit.scenario.scenarios.adaptive.selectors", diff --git a/pyrit/scenario/scenarios/adaptive/adaptive_scenario.py b/pyrit/scenario/scenarios/adaptive/adaptive_scenario.py index 16965aab65..c421f2974d 100644 --- a/pyrit/scenario/scenarios/adaptive/adaptive_scenario.py +++ b/pyrit/scenario/scenarios/adaptive/adaptive_scenario.py @@ -23,10 +23,16 @@ from pyrit.common.utils import to_sha256 from pyrit.executor.attack import AttackScoringConfig from pyrit.models import ( + AtomicAttackEvaluationIdentifier, + AtomicAttackIdentifier, + ScenarioAdaptiveRunSizeDetails, ScenarioRunSizeComponent, - ScenarioRunSizeEstimate, ) -from pyrit.models.identifiers import compute_inner_attack_eval_hash +from pyrit.models.catalog import ( + ScenarioDefaultRunSizeEstimate, + ScenarioRunSizeEstimateStatus, + ScenarioRunSizeFactor, +) from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.attack_technique import AttackTechnique from pyrit.scenario.core.matrix_atomic_attack_builder import build_baseline_atomic_attack @@ -34,6 +40,7 @@ from pyrit.scenario.core.scenario_target_defaults import get_default_adversarial_target from pyrit.scenario.scenarios.adaptive.dispatcher import AdaptiveTechniqueDispatcher, TechniqueBundle from pyrit.scenario.scenarios.adaptive.selectors import EpsilonGreedyTechniqueSelector, TechniqueSelector +from pyrit.scenario.scenarios.adaptive.technique_identity import AdaptiveTechniqueIdentifier if TYPE_CHECKING: from pyrit.models import AttackSeedGroup @@ -201,21 +208,30 @@ async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list return atomic_attacks - async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: + async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: """ Estimate compatible persisted envelopes, excluding adaptive inner attempts. Returns: - ScenarioRunSizeEstimate: The adaptive outer-envelope estimate. + ScenarioDefaultRunSizeEstimate: The adaptive outer-envelope estimate. + + Raises: + ValueError: If ``max_attempts_per_objective`` is less than one. """ selected_groups, datasets = await self._resolve_dataset_groups_for_estimate_async() selected_count = sum(len(groups) for groups in selected_groups.values()) max_attempts = int(self.params.get("max_attempts_per_objective", 3)) + if max_attempts < 1: + raise ValueError(f"max_attempts_per_objective must be >= 1, got {max_attempts}") + selected_candidate_count = len(self._scenario_techniques) + selected_attempt_bound = min(selected_candidate_count, max_attempts) + baseline_count = selected_count if self._include_baseline else 0 baseline_components = ( [ ScenarioRunSizeComponent( label="Baseline", count=selected_count, + factors=[ScenarioRunSizeFactor(label="objectives", count=selected_count)], is_baseline=True, ) ] @@ -226,24 +242,35 @@ async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: components = [ *baseline_components, ScenarioRunSizeComponent( - label="Adaptive attack-envelope candidates", + label="Adaptive objectives", count=selected_count, + factors=[ScenarioRunSizeFactor(label="objectives", count=selected_count)], ), ] - return ScenarioRunSizeEstimate( - minimum_attack_count=sum(component.count for component in baseline_components), - maximum_attack_count=sum(component.count for component in components), + return ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Conditional, + minimum_attack_count=baseline_count, + maximum_attack_count=baseline_count + selected_count, components=components, datasets=datasets, + adaptive_details=ScenarioAdaptiveRunSizeDetails( + objective_count=selected_count, + selected_candidate_technique_count=selected_candidate_count, + candidate_technique_count=selected_candidate_count, + max_attempts_per_objective=max_attempts, + techniques_per_objective_upper_bound=selected_attempt_bound, + technique_attempt_count_upper_bound=selected_count * selected_attempt_bound, + ), note=( - "The authoritative total depends on which selected techniques are compatible with the " - f"configured objective target and each seed group. Up to {max_attempts} inner attempts per " - "envelope and retries are excluded." + "The planned-attack total depends on which selected techniques are compatible with the " + f"configured objective target. Up to {selected_attempt_bound} selected technique attempts " + "may run per adaptive objective; retries are excluded." ), ) assert self._objective_target is not None techniques = self._build_techniques_dict(objective_target=self._objective_target) + candidate_count = len(techniques) dispatcher = AdaptiveTechniqueDispatcher( objective_target=self._objective_target, techniques=techniques, @@ -261,34 +288,52 @@ async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: components = [ *baseline_components, ScenarioRunSizeComponent( - label="Adaptive attack envelopes", + label="Adaptive objectives", count=compatible_group_count, + factors=[ScenarioRunSizeFactor(label="compatible objectives", count=compatible_group_count)], ), ] - estimated_attack_count = ( - None if self._estimate_has_binding_size_cap else sum(component.count for component in components) + status = ( + ScenarioRunSizeEstimateStatus.Conditional + if self._estimate_has_binding_size_cap + else ScenarioRunSizeEstimateStatus.Exact + ) + adaptive_objective_bound = ( + selected_count if status is ScenarioRunSizeEstimateStatus.Conditional else compatible_group_count ) - minimum_attack_count = None - maximum_attack_count = None + total_attack_count = ( + None + if status is ScenarioRunSizeEstimateStatus.Conditional + else sum(component.count for component in components) + ) + minimum_attack_count = ( + baseline_count if status is ScenarioRunSizeEstimateStatus.Conditional and baseline_count > 0 else None + ) + maximum_attack_count = ( + baseline_count + selected_count if status is ScenarioRunSizeEstimateStatus.Conditional else None + ) + technique_attempt_bound = min(candidate_count, max_attempts) note = ( - f"Each planned unit is one persisted adaptive envelope. Up to {max_attempts} selected technique " - "attempts may run inside that unit; inner attempts and retries are excluded." + f"Each compatible adaptive objective is one planned attack. Up to {technique_attempt_bound} selected " + "technique attempts may run for that objective; retries are excluded." ) - if estimated_attack_count is None: - baseline_count = sum(component.count for component in baseline_components) - minimum_attack_count = baseline_count - maximum_attack_count = baseline_count + selected_count - note += ( - " A binding randomized dataset cap may select a different compatibility mix at launch. " - "The range covers the baseline-only minimum through one compatible adaptive envelope per " - "selected seed group." - ) - return ScenarioRunSizeEstimate( - estimated_attack_count=estimated_attack_count, + if status is ScenarioRunSizeEstimateStatus.Conditional: + note += " A binding randomized dataset cap may select a different compatibility mix at launch." + return ScenarioDefaultRunSizeEstimate( + status=status, + total_attack_count=total_attack_count, minimum_attack_count=minimum_attack_count, maximum_attack_count=maximum_attack_count, components=components, datasets=datasets, + adaptive_details=ScenarioAdaptiveRunSizeDetails( + objective_count=adaptive_objective_bound, + selected_candidate_technique_count=selected_candidate_count, + candidate_technique_count=candidate_count, + max_attempts_per_objective=max_attempts, + techniques_per_objective_upper_bound=technique_attempt_bound, + technique_attempt_count_upper_bound=adaptive_objective_bound * technique_attempt_bound, + ), note=note, ) @@ -298,19 +343,17 @@ def _build_techniques_dict( objective_target: PromptTarget, ) -> dict[str, TechniqueBundle]: """ - Resolve selected techniques into a ``{eval_hash: TechniqueBundle}`` map. + Resolve selected techniques into a ``{adaptive_id: TechniqueBundle}`` map. Each bundle carries the inner attack technique along with the factory's ``seed_technique`` and ``adversarial_chat`` so the dispatcher can reproduce the static ``AtomicAttack`` execution path per attempt. - Technique keys are eval hashes derived from the inner attack technique's - identifier (run through ``AtomicAttackEvaluationIdentifier`` so seeds, - scorers, and operational target params are excluded). The same hash is - auto-stamped on every persisted ``AttackResultEntry.atomic_attack_identifier`` - by the executor, which lets the selector aggregate historical success - rates by behavioral configuration via - ``MemoryInterface.get_attack_results(atomic_attack_eval_hashes=...)``. + Technique keys join the canonical factory identifier hash with the full + ``AttackTechnique`` eval hash. Factory identity keeps distinct registered + configurations as separate arms even when they share an inner attack; + technique eval identity links those arms to normal-scenario history. + The dispatcher persists this joined identity on each child result. For factories whose attack class narrows ``attack_scoring_config`` to a specific subtype (e.g. ``TAPAttackScoringConfig`` for TAP), this method @@ -320,7 +363,7 @@ def _build_techniques_dict( are dropped with a warning so the rest of the pool continues to run. Returns: - dict[str, TechniqueBundle]: Mapping from technique eval hash to its + dict[str, TechniqueBundle]: Mapping from joined Adaptive identity to its bundle, in the order selected techniques were resolved. Raises: @@ -356,11 +399,17 @@ def _build_techniques_dict( skipped_incompatible[technique_name] = str(exc) logger.warning(f"Skipping technique '{technique_name}': {type(exc).__name__}: {exc}") continue - eval_hash = compute_inner_attack_eval_hash(attack=technique.attack) + technique_eval_hash = AtomicAttackEvaluationIdentifier( + AtomicAttackIdentifier.build(technique_identifier=technique.get_identifier()) + ).eval_hash + technique_identifier = AdaptiveTechniqueIdentifier( + factory_hash=factory.get_identifier().hash, + technique_eval_hash=technique_eval_hash, + ).serialize() adversarial_chat = factory.adversarial_chat if adversarial_chat is None and factory.uses_adversarial: adversarial_chat = get_default_adversarial_target() - techniques[eval_hash] = TechniqueBundle( + techniques[technique_identifier] = TechniqueBundle( attack=technique.attack, name=technique_name, seed_technique=technique.seed_technique, diff --git a/pyrit/scenario/scenarios/adaptive/dispatcher.py b/pyrit/scenario/scenarios/adaptive/dispatcher.py index ba4bc31f18..bafe0509b0 100644 --- a/pyrit/scenario/scenarios/adaptive/dispatcher.py +++ b/pyrit/scenario/scenarios/adaptive/dispatcher.py @@ -12,15 +12,10 @@ hands them to the scenario base for execution. The returned attack is a plain ``SequentialAttack`` with -``SequenceCompletionPolicy.FIRST_SUCCESS``. The per-attempt dispatch trail -(which technique ran, with what outcome, in what order) is not stamped onto -the envelope — every child ``AttackResult`` in -``SequentialAttackResult.child_attack_results`` already carries its own -``outcome`` and its own ``atomic_attack_identifier.eval_hash``. Callers that -want a human-readable technique label per child read it directly from the -child via ``child.get_attack_strategy_identifier().unique_name`` (the -executor auto-stamps ``class_name`` and ``unique_name`` on every persisted -row), so there is no separate ``{eval_hash: name}`` map to consult. +``SequenceCompletionPolicy.FIRST_SUCCESS``. Each child result carries the +stable registered-factory identity and friendly technique name in labels so +different configured techniques remain attributable even when their inner +attack implementations are execution-equivalent. """ from __future__ import annotations @@ -50,6 +45,12 @@ ADAPTIVE_ATTEMPT_LABEL: str = "_adaptive_attempt" """1-based attempt index within the per-objective loop.""" +ADAPTIVE_TECHNIQUE_ID_LABEL: str = "_adaptive_technique_id" +"""Joined registered-factory and behavioral-history identity for the selected arm.""" + +ADAPTIVE_TECHNIQUE_NAME_LABEL: str = "_adaptive_technique_name" +"""Registered technique name for human-readable result attribution.""" + @dataclass(frozen=True) class TechniqueBundle: @@ -58,16 +59,8 @@ class TechniqueBundle: Carries the inner attack strategy alongside the factory-supplied ``seed_technique`` (if any) and ``adversarial_chat`` (required when the - seed_technique contains a simulated-conversation config). ``name`` is the - factory-registration key; the dispatcher does not consume it, but it is - convenient for diagnostics and is preserved here so callers/tests can - cross-check which factory each bundle came from. - - Notebook/report code that wants a human-readable label for a persisted - child ``AttackResult`` should read it from the child itself via - ``child.get_attack_strategy_identifier()`` — the executor already stamps - ``class_name`` and ``unique_name`` on every row, so there is no need to - publish a separate ``{eval_hash: name}`` map. + seed technique contains a simulated-conversation config). ``name`` is the + factory-registration key and is persisted on every selected child result. """ attack: AttackStrategy[Any, AttackResult] @@ -110,7 +103,7 @@ def __init__( Args: objective_target (PromptTarget): The target inner attacks run against. techniques (dict[str, TechniqueBundle]): Mapping from - technique eval hash to its bundle. Must be non-empty. + joined Adaptive technique identity to its bundle. Must be non-empty. selector (TechniqueSelector): Stateless technique selector. objective_scorer (TrueFalseScorer | None): Scorer forwarded to inner attacks that generate simulated conversations. @@ -137,14 +130,14 @@ def __init__( def compatible_techniques(self, *, seed_group: AttackSeedGroup) -> list[str]: """ - Return technique hashes whose ``seed_technique`` is compatible with ``seed_group``. + Return technique identifiers whose ``seed_technique`` is compatible with ``seed_group``. Techniques with no ``seed_technique`` are universally compatible. Used by ``AdaptiveScenario`` to drop seed groups with no usable techniques before building atomic attacks. Returns: - list[str]: Technique eval hashes in declaration order. + list[str]: Joined Adaptive technique identities in declaration order. """ return [ name @@ -179,12 +172,8 @@ async def build_attack_async( technique map. Returns: - SequentialAttack: The ready-to-run attack. Each child's - identity is captured by its own - ``atomic_attack_identifier.eval_hash`` after execution; - callers wanting the friendly technique name read it - directly from the child via - ``child.get_attack_strategy_identifier().unique_name``. + SequentialAttack: The ready-to-run attack. Each child carries its + canonical factory identity and registered name in result labels. Raises: ValueError: If ``seed_group.objective`` is not initialized, @@ -202,7 +191,7 @@ async def build_attack_async( f"(objective={seed_group.objective.value!r})." ) - chosen_hashes = await self._selector.select_async( + chosen_identifiers = await self._selector.select_async( technique_identifiers=compatible, objective=seed_group.objective.value, num_top_techniques=self._max_attempts, @@ -210,7 +199,7 @@ async def build_attack_async( ) child_attacks: list[SequentialChildAttack] = [] - for attempt_idx, chosen in enumerate(chosen_hashes): + for attempt_idx, chosen in enumerate(chosen_identifiers): bundle = self._techniques[chosen] execution_group = ( seed_group.with_technique(technique=bundle.seed_technique) @@ -223,7 +212,11 @@ async def build_attack_async( seed_group=execution_group, adversarial_chat=bundle.adversarial_chat, objective_scorer=self._objective_scorer, - memory_labels={ADAPTIVE_ATTEMPT_LABEL: str(attempt_idx + 1)}, + memory_labels={ + ADAPTIVE_ATTEMPT_LABEL: str(attempt_idx + 1), + ADAPTIVE_TECHNIQUE_ID_LABEL: chosen, + ADAPTIVE_TECHNIQUE_NAME_LABEL: bundle.name, + }, ) ) diff --git a/pyrit/scenario/scenarios/adaptive/selectors/epsilon_greedy.py b/pyrit/scenario/scenarios/adaptive/selectors/epsilon_greedy.py index a415182733..3fd2b000e9 100644 --- a/pyrit/scenario/scenarios/adaptive/selectors/epsilon_greedy.py +++ b/pyrit/scenario/scenarios/adaptive/selectors/epsilon_greedy.py @@ -11,8 +11,10 @@ import struct from typing import TYPE_CHECKING -from pyrit.analytics.technique_analysis import compute_technique_stats +from pyrit.analytics.technique_analysis import compute_labeled_technique_stats +from pyrit.scenario.scenarios.adaptive.dispatcher import ADAPTIVE_TECHNIQUE_ID_LABEL from pyrit.scenario.scenarios.adaptive.selectors.technique_selector import SelectorScope +from pyrit.scenario.scenarios.adaptive.technique_identity import get_history_eval_hash if TYPE_CHECKING: from collections.abc import Sequence @@ -129,8 +131,13 @@ async def select_async( rng = _derive_rng(self._seed, decision_key) effective_run_id = scenario_result_id if self._scope.current_run_only else None - stats = compute_technique_stats( - technique_eval_hashes=technique_list, + stats = compute_labeled_technique_stats( + technique_identifiers=technique_list, + label_name=ADAPTIVE_TECHNIQUE_ID_LABEL, + technique_eval_hashes_by_identifier={ + technique_identifier: get_history_eval_hash(technique_identifier=technique_identifier) + for technique_identifier in technique_list + }, scenario_result_id=effective_run_id, targeted_harm_categories=self._scope.targeted_harm_categories, ) diff --git a/pyrit/scenario/scenarios/adaptive/selectors/technique_selector.py b/pyrit/scenario/scenarios/adaptive/selectors/technique_selector.py index 0161e4923b..4ae383df55 100644 --- a/pyrit/scenario/scenarios/adaptive/selectors/technique_selector.py +++ b/pyrit/scenario/scenarios/adaptive/selectors/technique_selector.py @@ -20,7 +20,7 @@ class SelectorScope: All fields default to "no restriction"; combine fields to narrow the scope (e.g. current run only, same harm category). Filter values flow - through ``compute_technique_stats`` to + through labeled technique statistics to ``MemoryInterface.get_attack_results``. The scope is held by the selector at construction time. The per-call @@ -28,10 +28,9 @@ class SelectorScope: to memory only when ``current_run_only`` is set; otherwise the selector queries across all runs. - Per-technique disambiguation uses ``atomic_attack_identifier.eval_hash`` - (auto-stamped on every persisted attack result), which already encodes - the attack class plus its behavior-relevant params. Class-based - narrowing is therefore unnecessary at this layer. + Per-technique disambiguation uses the canonical registered factory + identifier persisted on every Adaptive child result. This keeps distinct + registered configurations separate even when they share an attack class. """ current_run_only: bool = False @@ -86,8 +85,8 @@ async def select_async( Return techniques in priority order (try first, try second, …). Args: - technique_identifiers (Sequence[str]): Available technique eval - hashes. + technique_identifiers (Sequence[str]): Available stable Adaptive + technique identities. objective (str): The objective text for this selection. num_top_techniques (int): Max techniques to return. Defaults to 1. scenario_result_id (str | None): The current scenario run ID, @@ -96,7 +95,7 @@ async def select_async( ``current_run_only=True``. Returns: - Sequence[str]: Up to ``num_top_techniques`` technique eval hashes + Sequence[str]: Up to ``num_top_techniques`` technique identities in priority order. Fewer if not enough techniques are available. """ diff --git a/pyrit/scenario/scenarios/adaptive/technique_identity.py b/pyrit/scenario/scenarios/adaptive/technique_identity.py new file mode 100644 index 0000000000..950d7358b7 --- /dev/null +++ b/pyrit/scenario/scenarios/adaptive/technique_identity.py @@ -0,0 +1,72 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Stable identity carried by Adaptive selector arms.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import ClassVar + + +@dataclass(frozen=True) +class AdaptiveTechniqueIdentifier: + """ + Join registered factory identity with cross-scenario behavioral identity. + + The factory hash keeps separately registered configured techniques as + distinct selector arms. The technique eval hash links each arm to normal + scenario results persisted from the same ``AttackTechnique`` behavior. + """ + + factory_hash: str + technique_eval_hash: str + + _PREFIX: ClassVar[str] = "adaptive-v1" + _SEPARATOR: ClassVar[str] = ":" + + def serialize(self) -> str: + """ + Serialize the identifier for selector keys and persisted labels. + + Returns: + str: Versioned identifier containing both canonical hashes. + + Raises: + ValueError: If either hash is empty or contains the field separator. + """ + if self._SEPARATOR in self.factory_hash or self._SEPARATOR in self.technique_eval_hash: + raise ValueError("Adaptive technique identity hashes cannot contain ':'") + if not self.factory_hash or not self.technique_eval_hash: + raise ValueError("Adaptive technique identity hashes cannot be empty") + return self._SEPARATOR.join((self._PREFIX, self.factory_hash, self.technique_eval_hash)) + + @classmethod + def parse(cls, value: str) -> AdaptiveTechniqueIdentifier | None: + """ + Parse a serialized Adaptive identifier. + + Unknown selector identifiers remain valid for custom selectors and + legacy tests, so malformed or unversioned values return ``None``. + + Returns: + AdaptiveTechniqueIdentifier | None: Parsed identity when recognized. + """ + parts = value.split(cls._SEPARATOR) + if len(parts) != 3 or parts[0] != cls._PREFIX or not parts[1] or not parts[2]: + return None + return cls(factory_hash=parts[1], technique_eval_hash=parts[2]) + + +def get_history_eval_hash(*, technique_identifier: str) -> str: + """ + Return the normal-scenario eval hash associated with a selector arm. + + Unversioned identifiers fall back to themselves for backward compatibility + with custom selectors and callers that already pass eval hashes. + + Returns: + str: Behavioral eval hash used for historical result lookup. + """ + parsed = AdaptiveTechniqueIdentifier.parse(technique_identifier) + return parsed.technique_eval_hash if parsed is not None else technique_identifier diff --git a/pyrit/scenario/scenarios/adaptive/text_adaptive.py b/pyrit/scenario/scenarios/adaptive/text_adaptive.py index ad1c468312..256e94ec22 100644 --- a/pyrit/scenario/scenarios/adaptive/text_adaptive.py +++ b/pyrit/scenario/scenarios/adaptive/text_adaptive.py @@ -126,7 +126,10 @@ def additional_parameters(cls) -> list[Parameter]: return [ Parameter( name="max_attempts_per_objective", - description="Max techniques tried per objective. Defaults to 3.", + description=( + "Maximum different compatible techniques Adaptive may try for one objective, stopping after " + "the first success. This is separate from retries." + ), param_type=int, default=3, ), diff --git a/pyrit/scenario/scenarios/airt/jailbreak.py b/pyrit/scenario/scenarios/airt/jailbreak.py index ac94993d52..e9b956b110 100644 --- a/pyrit/scenario/scenarios/airt/jailbreak.py +++ b/pyrit/scenario/scenarios/airt/jailbreak.py @@ -16,7 +16,12 @@ AttackTechniqueSeedGroup, Parameter, ScenarioRunSizeComponent, - ScenarioRunSizeEstimate, + ScenarioRunSizeEstimateCondition, +) +from pyrit.models.catalog import ( + ScenarioDefaultRunSizeEstimate, + ScenarioRunSizeEstimateStatus, + ScenarioRunSizeFactor, ) from pyrit.prompt_target import CapabilityName from pyrit.registry.components.attack_technique_registry import AttackTechniqueRegistry @@ -316,12 +321,12 @@ def _build_initial_scenario_metadata(self) -> dict[str, Any]: metadata[_JAILBREAK_TEMPLATES_METADATA_KEY] = list(self._resolved_jailbreaks) return metadata - async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: + async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: """ Estimate the template and attempt axes, preserving the target capability caveat. Returns: - ScenarioRunSizeEstimate: Conditional target-aware estimate. + ScenarioDefaultRunSizeEstimate: Conditional target-aware estimate. Raises: ValueError: If native system-prompt delivery is the only selected @@ -353,6 +358,7 @@ async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: ScenarioRunSizeComponent( label="Baseline", count=seed_group_count, + factors=[ScenarioRunSizeFactor(label="selected logical seed groups", count=seed_group_count)], is_baseline=True, ) ) @@ -360,6 +366,12 @@ async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: ScenarioRunSizeComponent( label="Inline jailbreak delivery", count=seed_group_count * template_count * attempt_count * converter_count, + factors=[ + ScenarioRunSizeFactor(label="selected logical seed groups", count=seed_group_count), + ScenarioRunSizeFactor(label="jailbreak templates", count=template_count), + ScenarioRunSizeFactor(label="attempts", count=attempt_count), + ScenarioRunSizeFactor(label="inline delivery techniques", count=converter_count), + ], note=( "Each planned unit is one template, one selected delivery technique, and one logical seed group. " "num_jailbreaks selects templates; it is not a persisted result or attempt count." @@ -371,6 +383,12 @@ async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: ScenarioRunSizeComponent( label="Native system-prompt jailbreak delivery", count=seed_group_count * template_count * attempt_count, + factors=[ + ScenarioRunSizeFactor(label="selected logical seed groups", count=seed_group_count), + ScenarioRunSizeFactor(label="jailbreak templates", count=template_count), + ScenarioRunSizeFactor(label="attempts", count=attempt_count), + ], + condition=ScenarioRunSizeEstimateCondition.TargetCapabilities, note=( "The selected objective target supports native system-prompt delivery." if system_delivery_supported is True @@ -396,10 +414,12 @@ async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: f"{converter_count} selected target-agnostic technique(s) x {attempt_count} configured attempt(s) " f"= {seed_group_count * template_count * attempt_count * converter_count} planned unit(s)." ) - estimated_attack_count = ( - None if system_delivery_selected and system_delivery_supported is None else planned_count + status = ( + ScenarioRunSizeEstimateStatus.Conditional + if system_delivery_selected and system_delivery_supported is None + else ScenarioRunSizeEstimateStatus.Exact ) - if estimated_attack_count is None: + if status is ScenarioRunSizeEstimateStatus.Conditional: capability_note = ( " The selected technique requires native system-prompt delivery; incompatible targets cannot run it." if converter_count == 0 @@ -422,10 +442,18 @@ async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: effective_parameters["jailbreak_names"] = list(jailbreak_names) else: effective_parameters["num_jailbreaks"] = template_count - return ScenarioRunSizeEstimate( - estimated_attack_count=estimated_attack_count, - minimum_attack_count=minimum_planned_count if estimated_attack_count is None else None, - maximum_attack_count=planned_count if estimated_attack_count is None else None, + return ScenarioDefaultRunSizeEstimate( + status=status, + total_attack_count=planned_count if status is ScenarioRunSizeEstimateStatus.Exact else None, + minimum_attack_count=( + minimum_planned_count if status is ScenarioRunSizeEstimateStatus.Conditional else None + ), + maximum_attack_count=planned_count if status is ScenarioRunSizeEstimateStatus.Conditional else None, + condition=( + ScenarioRunSizeEstimateCondition.TargetCapabilities + if status is ScenarioRunSizeEstimateStatus.Conditional + else None + ), components=components, datasets=datasets, effective_parameters=effective_parameters, diff --git a/pyrit/scenario/scenarios/airt/psychosocial.py b/pyrit/scenario/scenarios/airt/psychosocial.py index af6b28e193..338931f643 100644 --- a/pyrit/scenario/scenarios/airt/psychosocial.py +++ b/pyrit/scenario/scenarios/airt/psychosocial.py @@ -6,7 +6,7 @@ import logging import pathlib from dataclasses import dataclass -from typing import TYPE_CHECKING, cast +from typing import TYPE_CHECKING, ClassVar, Literal, cast from pyrit.common import apply_defaults from pyrit.common.path import DATASETS_PATH @@ -342,6 +342,9 @@ class Psychosocial(Scenario): """ VERSION: int = 3 + DATASET_SIZE_LIMIT_OVERRIDE_SCOPE: ClassVar[Literal["per_dataset", "combined", "unsupported"] | None] = ( + "per_dataset" + ) @classmethod def additional_parameters(cls) -> list[Parameter]: diff --git a/pyrit/scenario/scenarios/garak/web_injection.py b/pyrit/scenario/scenarios/garak/web_injection.py index ddf7216d4b..5b6bf4e997 100644 --- a/pyrit/scenario/scenarios/garak/web_injection.py +++ b/pyrit/scenario/scenarios/garak/web_injection.py @@ -6,7 +6,7 @@ import asyncio import logging import random -from typing import TYPE_CHECKING, ClassVar, cast +from typing import TYPE_CHECKING, ClassVar, Literal, cast from pyrit.common import apply_defaults from pyrit.executor.attack.core.attack_config import AttackScoringConfig @@ -97,6 +97,9 @@ class WebInjection(Scenario): VERSION: int = 1 BASELINE_ATTACK_POLICY: ClassVar[BaselineAttackPolicy] = BaselineAttackPolicy.Enabled + DATASET_SIZE_LIMIT_OVERRIDE_SCOPE: ClassVar[Literal["per_dataset", "combined", "unsupported"] | None] = ( + "unsupported" + ) # Local ``.prompt`` datasets under datasets/seed_datasets/local/garak. DATASET_EXAMPLE_DOMAINS: ClassVar[str] = "garak_example_domains_xss" diff --git a/tests/unit/analytics/test_technique_analysis.py b/tests/unit/analytics/test_technique_analysis.py index 04b1d94890..c31a1cd0ea 100644 --- a/tests/unit/analytics/test_technique_analysis.py +++ b/tests/unit/analytics/test_technique_analysis.py @@ -2,15 +2,17 @@ # Licensed under the MIT license. from unittest.mock import MagicMock, patch +from uuid import uuid4 import pytest -from pyrit.analytics.technique_analysis import compute_technique_stats +from pyrit.analytics.technique_analysis import compute_labeled_technique_stats, compute_technique_stats from pyrit.models import AttackOutcome def _make_result(*, eval_hash: str | None, outcome: AttackOutcome) -> MagicMock: r = MagicMock() + r.attack_result_id = str(uuid4()) if eval_hash is None: r.atomic_attack_identifier = None else: @@ -18,9 +20,16 @@ def _make_result(*, eval_hash: str | None, outcome: AttackOutcome) -> MagicMock: identifier.eval_hash = eval_hash r.atomic_attack_identifier = identifier r.outcome = outcome + r.labels = {} return r +def _make_labeled_result(*, label_name: str, technique_identifier: str, outcome: AttackOutcome) -> MagicMock: + result = _make_result(eval_hash=None, outcome=outcome) + result.labels = {label_name: technique_identifier} + return result + + @pytest.fixture(autouse=True) def _patch_memory(): mock_memory = MagicMock() @@ -144,3 +153,77 @@ def test_injected_memory_bypasses_central_memory(self, _patch_memory): injected.get_attack_results.assert_called_once() _patch_memory.get_attack_results.assert_not_called() assert stats["a"].successes == 1 + + +class TestComputeLabeledTechniqueStats: + def test_counts_distinct_labeled_techniques(self, _patch_memory): + label_name = "_adaptive_technique_id" + _patch_memory.get_attack_results.return_value = [ + _make_labeled_result( + label_name=label_name, + technique_identifier="role-play-movie", + outcome=AttackOutcome.SUCCESS, + ), + _make_labeled_result( + label_name=label_name, + technique_identifier="role-play-video", + outcome=AttackOutcome.FAILURE, + ), + ] + + stats = compute_labeled_technique_stats( + technique_identifiers=["role-play-movie", "role-play-video"], + label_name=label_name, + ) + + assert stats["role-play-movie"].successes == 1 + assert stats["role-play-video"].failures == 1 + assert _patch_memory.get_attack_results.call_args.kwargs["labels"] == { + label_name: ["role-play-movie", "role-play-video"] + } + + def test_unlabeled_and_unrequested_results_are_ignored(self, _patch_memory): + label_name = "_adaptive_technique_id" + unlabeled = _make_result(eval_hash="shared", outcome=AttackOutcome.SUCCESS) + _patch_memory.get_attack_results.return_value = [ + unlabeled, + _make_labeled_result( + label_name=label_name, + technique_identifier="other", + outcome=AttackOutcome.SUCCESS, + ), + ] + + stats = compute_labeled_technique_stats( + technique_identifiers=["requested"], + label_name=label_name, + ) + + assert stats == {} + + def test_merges_labeled_and_normal_scenario_history_without_double_counting(self, _patch_memory): + label_name = "_adaptive_technique_id" + labeled = _make_labeled_result( + label_name=label_name, + technique_identifier="factory-arm", + outcome=AttackOutcome.SUCCESS, + ) + labeled.atomic_attack_identifier = MagicMock(eval_hash="inner-attack-hash") + normal = _make_result(eval_hash="full-technique-hash", outcome=AttackOutcome.FAILURE) + _patch_memory.get_attack_results.side_effect = [ + [labeled], + [labeled, normal], + ] + + stats = compute_labeled_technique_stats( + technique_identifiers=["factory-arm"], + label_name=label_name, + technique_eval_hashes_by_identifier={"factory-arm": "full-technique-hash"}, + ) + + assert stats["factory-arm"].successes == 1 + assert stats["factory-arm"].failures == 1 + assert stats["factory-arm"].total_decided == 2 + assert _patch_memory.get_attack_results.call_args_list[1].kwargs["atomic_attack_eval_hashes"] == [ + "full-technique-hash" + ] diff --git a/tests/unit/backend/test_scenario_run_service.py b/tests/unit/backend/test_scenario_run_service.py index c2a5b75fbb..8e13888613 100644 --- a/tests/unit/backend/test_scenario_run_service.py +++ b/tests/unit/backend/test_scenario_run_service.py @@ -45,7 +45,11 @@ config_hash, ) from pyrit.models.catalog.scenario import RunScenarioRequest -from pyrit.scenario.core import DatasetAttackConfiguration, DatasetConfiguration +from pyrit.scenario.core import ( + CompoundDatasetAttackConfiguration, + DatasetAttackConfiguration, + DatasetConfiguration, +) from pyrit.scenario.core.scenario_technique import ScenarioTechnique from unit.mocks import make_scenario_result @@ -585,6 +589,81 @@ class _MarkerDatasetConfiguration(DatasetConfiguration): assert built_config.dataset_names == ["only_this"] assert built_config.max_dataset_size is None + async def test_start_run_dataset_names_rebuilds_homogeneous_compound(self, mock_all_registries) -> None: + """Compound per-dataset defaults support exact selected-name overrides.""" + default_config = CompoundDatasetAttackConfiguration.per_dataset( + dataset_names=["airt_hate", "airt_fairness"], + max_dataset_size=4, + ) + scenario_instance = mock_all_registries["scenario_instance"] + scenario_instance._default_dataset_config = default_config + + service = ScenarioRunService() + await service.start_run_async(request=_make_request(dataset_names=["airt_fairness"])) + + init_call = mock_all_registries["scenario_registry"].create_and_initialize_async.await_args + built_config = init_call.kwargs["dataset_config"] + assert isinstance(built_config, CompoundDatasetAttackConfiguration) + assert built_config.dataset_names == ["airt_fairness"] + assert [child.max_dataset_size for child in built_config._configurations] == [4] + assert default_config.dataset_names == ["airt_hate", "airt_fairness"] + + async def test_start_run_max_dataset_size_updates_each_default_compound_child(self, mock_all_registries) -> None: + """An unchanged default selection keeps compound caps per dataset.""" + default_config = CompoundDatasetAttackConfiguration.per_dataset( + dataset_names=["airt_hate", "airt_fairness"], + max_dataset_size=4, + ) + scenario_instance = mock_all_registries["scenario_instance"] + scenario_instance._default_dataset_config = default_config + + service = ScenarioRunService() + await service.start_run_async(request=_make_request(max_dataset_size=2)) + + init_call = mock_all_registries["scenario_registry"].create_and_initialize_async.await_args + built_config = init_call.kwargs["dataset_config"] + assert isinstance(built_config, CompoundDatasetAttackConfiguration) + assert built_config is default_config + assert built_config.dataset_names == ["airt_hate", "airt_fairness"] + assert [child.max_dataset_size for child in built_config._configurations] == [2, 2] + + async def test_start_run_non_name_overrides_preserve_shaped_compound_children(self, mock_all_registries) -> None: + """Size and filter overrides do not rebuild scenario-specific child configurations.""" + + class _ShapedDatasetConfiguration(DatasetAttackConfiguration): + pass + + default_config = CompoundDatasetAttackConfiguration( + configurations=[ + _ShapedDatasetConfiguration(dataset_names=["d1"], max_dataset_size=4), + _ShapedDatasetConfiguration(dataset_names=["d2"], max_dataset_size=4), + ], + ) + scenario_instance = mock_all_registries["scenario_instance"] + scenario_instance._default_dataset_config = default_config + + service = ScenarioRunService() + await service.start_run_async( + request=_make_request( + dataset_names=["d1", "d2"], + max_dataset_size=2, + 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"] + assert built_config is default_config + assert [type(child) for child in built_config._configurations] == [ + _ShapedDatasetConfiguration, + _ShapedDatasetConfiguration, + ] + assert [child.max_dataset_size for child in built_config._configurations] == [2, 2] + assert [child.filters for child in built_config._configurations] == [ + {"harm_categories": ["cyber"]}, + {"harm_categories": ["cyber"]}, + ] + async def test_start_run_dataset_names_rejects_incompatible_subclass_constructor(self, mock_all_registries) -> None: """Reject overrides that cannot preserve scenario-specific dataset configuration.""" diff --git a/tests/unit/backend/test_scenario_service.py b/tests/unit/backend/test_scenario_service.py index ff657eb456..ceed2f32a8 100644 --- a/tests/unit/backend/test_scenario_service.py +++ b/tests/unit/backend/test_scenario_service.py @@ -7,6 +7,7 @@ import asyncio from collections import OrderedDict +from collections.abc import Awaitable from typing import Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -26,13 +27,18 @@ from pyrit.models import ( Parameter, ScenarioDatasetSizeCap, + ScenarioDatasetSizeLimit, ScenarioDatasetSummary, ScenarioRunSizeComponent, ScenarioRunSizeEstimate, ScenarioRunSizeEstimateRequest, ScenarioTechniqueSummary, ) -from pyrit.models.catalog.scenario import RegisteredScenario +from pyrit.models.catalog import ( + RegisteredScenario, + ScenarioDefaultRunSizeEstimate, + ScenarioRunSizeEstimateStatus, +) from pyrit.registry import ScenarioMetadata from pyrit.scenario.core import DatasetAttackConfiguration, ScenarioTechnique @@ -66,6 +72,15 @@ def clear_service_cache(): get_scenario_service.cache_clear() +def _initialize_test_service(service: ScenarioService) -> None: + """Initialize service state without binding the process-wide registry.""" + service._estimate_cache = OrderedDict() + service._estimate_tasks = OrderedDict() + service._estimate_task_lock = asyncio.Lock() + service._estimate_semaphore = asyncio.Semaphore(4) + service._configured_estimate_semaphore = asyncio.Semaphore(4) + + def _make_scenario_metadata( *, registry_name: str = "test.scenario", @@ -96,6 +111,7 @@ def _make_scenario_metadata( default_datasets: tuple[str, ...] = ("test_dataset",), baseline_policy: str = "enabled", include_baseline_by_default: bool = True, + dataset_size_limit: ScenarioDatasetSizeLimit | None = None, ) -> ScenarioMetadata: """Create a ScenarioMetadata instance for testing.""" return ScenarioMetadata( @@ -112,6 +128,7 @@ def _make_scenario_metadata( aggregate_technique_expansions=aggregate_technique_expansions, technique_summaries=technique_summaries, default_datasets=default_datasets, + dataset_size_limit=dataset_size_limit or ScenarioDatasetSizeLimit(), baseline_policy=baseline_policy, include_baseline_by_default=include_baseline_by_default, ) @@ -127,7 +144,7 @@ class TestScenarioServiceListScenarios: async def test_list_scenarios_returns_empty_when_no_scenarios(self) -> None: """Test that list returns empty list when no scenarios are registered.""" - with patch.object(ScenarioService, "__init__", lambda self: None): + with patch.object(ScenarioService, "__init__", _initialize_test_service): service = ScenarioService() service._registry = MagicMock() service._registry.get_all_registered_class_metadata.return_value = [] @@ -141,7 +158,7 @@ async def test_list_scenarios_returns_scenarios_from_registry(self) -> None: """Test that list returns scenarios from registry.""" metadata = _make_scenario_metadata() - with patch.object(ScenarioService, "__init__", lambda self: None): + with patch.object(ScenarioService, "__init__", _initialize_test_service): service = ScenarioService() service._registry = MagicMock() service._registry.get_all_registered_class_metadata.return_value = [metadata] @@ -161,6 +178,7 @@ async def test_list_scenarios_returns_scenarios_from_registry(self) -> None: assert result.items[0].technique_summaries[0].description == "Frames the objective as role play." assert result.items[0].technique_summaries[0].tags == ["default", "single_turn"] assert result.items[0].default_datasets == ["test_dataset"] + assert result.items[0].dataset_size_limit == ScenarioDatasetSizeLimit() assert result.items[0].baseline_policy == "enabled" assert result.items[0].include_baseline_by_default is True @@ -179,6 +197,24 @@ async def test_list_scenarios_can_return_metadata_without_waiting_for_estimates( assert result.items[0].default_run_size == ScenarioRunSizeEstimate.unavailable() service._registry.create_instance.assert_not_called() + async def test_list_scenarios_projects_dataset_size_limit_metadata(self) -> None: + """Catalog responses preserve structured scenario-owned limit semantics.""" + limit = ScenarioDatasetSizeLimit( + default_scope="per_dataset", + default_count=4, + override_scope="per_dataset", + ) + metadata = _make_scenario_metadata(dataset_size_limit=limit) + + with patch.object(ScenarioService, "__init__", _initialize_test_service): + service = ScenarioService() + service._registry = MagicMock() + service._registry.get_all_registered_class_metadata.return_value = [metadata] + + result = await service.list_scenarios_async() + + assert result.items[0].dataset_size_limit == limit + async def test_estimate_is_offloaded_and_cached(self) -> None: """Scenario-owned estimates run in a worker once and are reused by subsequent reads.""" metadata = _make_scenario_metadata() @@ -204,7 +240,7 @@ async def test_estimate_is_offloaded_and_cached(self) -> None: scenario = MagicMock() scenario.get_default_run_size_estimate_async = AsyncMock(return_value=estimate) - with patch.object(ScenarioService, "__init__", lambda self: None): + with patch.object(ScenarioService, "__init__", _initialize_test_service): service = ScenarioService() service._registry = MagicMock() service._registry.get_registered_class_metadata.return_value = metadata @@ -239,7 +275,7 @@ async def estimate_async() -> ScenarioRunSizeEstimate: scenario = MagicMock() scenario.get_default_run_size_estimate_async = AsyncMock(side_effect=estimate_async) - with patch.object(ScenarioService, "__init__", lambda self: None): + with patch.object(ScenarioService, "__init__", _initialize_test_service): service = ScenarioService() service._registry = MagicMock() service._registry.create_instance.return_value = scenario @@ -289,7 +325,7 @@ async def estimate_async() -> ScenarioRunSizeEstimate: scenario = MagicMock() scenario.get_default_run_size_estimate_async = AsyncMock(side_effect=estimate_async) - with patch.object(ScenarioService, "__init__", lambda self: None): + with patch.object(ScenarioService, "__init__", _initialize_test_service): service = ScenarioService() service._registry = MagicMock() service._registry.create_instance.return_value = scenario @@ -319,7 +355,7 @@ async def test_completed_stale_task_cannot_block_inflight_capacity(self) -> None scenario.get_default_run_size_estimate_async = AsyncMock(return_value=estimate) with ( - patch.object(ScenarioService, "__init__", lambda self: None), + patch.object(ScenarioService, "__init__", _initialize_test_service), patch("pyrit.backend.services.scenario_service._ESTIMATE_INFLIGHT_SIZE", 1), ): service = ScenarioService() @@ -351,7 +387,7 @@ async def test_one_failed_estimate_does_not_break_catalog(self) -> None: bad_scenario = MagicMock() bad_scenario.get_default_run_size_estimate_async = AsyncMock(side_effect=RuntimeError("dataset unavailable")) - with patch.object(ScenarioService, "__init__", lambda self: None): + with patch.object(ScenarioService, "__init__", _initialize_test_service): service = ScenarioService() service._registry = MagicMock() service._registry.get_all_registered_class_metadata.return_value = metadata @@ -363,6 +399,170 @@ async def test_one_failed_estimate_does_not_break_catalog(self) -> None: result = await service.list_scenarios_async() assert "RuntimeError" in result.items[1].default_run_size.note + async def test_catalog_estimates_use_bounded_parallelism(self) -> None: + """Catalog cards estimate concurrently without exceeding the configured bound.""" + metadata = [_make_scenario_metadata(registry_name=f"test.scenario_{index}") for index in range(6)] + estimate = ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Exact, + total_attack_count=1, + components=[ScenarioRunSizeComponent(label="Default sweep", count=1)], + ) + active = 0 + maximum_active = 0 + + async def estimate_async() -> ScenarioDefaultRunSizeEstimate: + nonlocal active, maximum_active + active += 1 + maximum_active = max(maximum_active, active) + await asyncio.sleep(0.01) + active -= 1 + return estimate + + scenario = MagicMock() + scenario.get_default_run_size_estimate_async = AsyncMock(side_effect=estimate_async) + service = ScenarioService() + service._registry = MagicMock() + service._registry.get_all_registered_class_metadata.return_value = metadata + service._registry.create_instance.return_value = scenario + service._estimate_semaphore = asyncio.Semaphore(2) + + result = await service.list_scenarios_async() + + assert maximum_active == 2 + assert all(item.default_run_size == estimate for item in result.items) + + async def test_catalog_queue_wait_does_not_start_execution_timeout(self) -> None: + """A queued catalog estimate starts its execution timeout only after acquiring capacity.""" + metadata = [_make_scenario_metadata(registry_name=f"test.scenario_{index}") for index in range(2)] + estimate = ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Exact, + total_attack_count=1, + components=[ScenarioRunSizeComponent(label="Default sweep", count=1)], + ) + first_estimate_started = asyncio.Event() + release_first_estimate = asyncio.Event() + second_timeout_started = asyncio.Event() + estimate_count = 0 + timeout_count = 0 + + async def estimate_async() -> ScenarioDefaultRunSizeEstimate: + nonlocal estimate_count + estimate_count += 1 + if estimate_count == 1: + first_estimate_started.set() + await release_first_estimate.wait() + return estimate + + async def wait_for_async( + awaitable: Awaitable[ScenarioDefaultRunSizeEstimate], *, timeout: float + ) -> ScenarioDefaultRunSizeEstimate: + nonlocal timeout_count + timeout_count += 1 + if timeout_count == 2: + second_timeout_started.set() + return await awaitable + + scenario = MagicMock() + scenario.get_default_run_size_estimate_async = AsyncMock(side_effect=estimate_async) + service = ScenarioService() + service._registry = MagicMock() + service._registry.get_all_registered_class_metadata.return_value = metadata + service._registry.create_instance.return_value = scenario + service._estimate_semaphore = asyncio.Semaphore(1) + + with patch("pyrit.backend.services.scenario_service.asyncio.wait_for", side_effect=wait_for_async): + catalog_task = asyncio.create_task(service.list_scenarios_async()) + await first_estimate_started.wait() + await asyncio.sleep(0) + assert not second_timeout_started.is_set() + + release_first_estimate.set() + result = await catalog_task + + assert second_timeout_started.is_set() + assert timeout_count == 2 + assert all(item.default_run_size == estimate for item in result.items) + + async def test_catalog_execution_timeout_is_unavailable_and_cached(self) -> None: + """A genuine estimate execution timeout is unavailable and reused from cache.""" + metadata = _make_scenario_metadata() + estimate_started = asyncio.Event() + estimate_cancelled = asyncio.Event() + block_estimate = asyncio.Event() + + async def slow_estimate_async() -> ScenarioDefaultRunSizeEstimate: + estimate_started.set() + try: + await block_estimate.wait() + except asyncio.CancelledError: + estimate_cancelled.set() + raise + raise AssertionError("The blocked estimate should be cancelled by its execution timeout.") + + scenario = MagicMock() + scenario.get_default_run_size_estimate_async = AsyncMock(side_effect=slow_estimate_async) + service = ScenarioService() + service._registry = MagicMock() + service._registry.create_instance.return_value = scenario + + with patch("pyrit.backend.services.scenario_service._DEFAULT_ESTIMATE_TIMEOUT_SECONDS", 0.01): + estimate_task = asyncio.create_task(service._get_default_run_size_estimate_async(metadata=metadata)) + await estimate_started.wait() + result = await estimate_task + cached = await service._get_default_run_size_estimate_async(metadata=metadata) + await asyncio.sleep(0) + + assert result.status is ScenarioRunSizeEstimateStatus.Unavailable + assert cached is result + assert estimate_cancelled.is_set() + service._registry.create_instance.assert_called_once_with(metadata.registry_name) + scenario.get_default_run_size_estimate_async.assert_awaited_once() + assert service._estimate_tasks == {} + + async def test_configured_estimate_does_not_wait_for_catalog_estimate(self) -> None: + """Interactive detail estimates use separate capacity from default catalog cards.""" + metadata = _make_scenario_metadata() + default_estimate = ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Exact, + total_attack_count=1, + components=[ScenarioRunSizeComponent(label="Default sweep", count=1)], + ) + configured_estimate = ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Exact, + total_attack_count=2, + components=[ScenarioRunSizeComponent(label="Configured sweep", count=2)], + ) + started = asyncio.Event() + release = asyncio.Event() + + async def default_estimate_async() -> ScenarioDefaultRunSizeEstimate: + started.set() + await release.wait() + return default_estimate + + scenario = MagicMock() + scenario.get_default_run_size_estimate_async = AsyncMock(side_effect=default_estimate_async) + service = ScenarioService() + service._registry = MagicMock() + service._registry.get_registered_class_metadata.return_value = metadata + service._registry.create_instance.return_value = scenario + service._estimate_semaphore = asyncio.Semaphore(1) + service._estimate_configured_run_size_async = AsyncMock(return_value=configured_estimate) + + catalog_task = asyncio.create_task(service._get_default_run_size_estimate_async(metadata=metadata)) + await started.wait() + interactive = await asyncio.wait_for( + service.estimate_scenario_run_size_async( + scenario_name=metadata.registry_name, + request=ScenarioRunSizeEstimateRequest(), + ), + timeout=0.2, + ) + release.set() + + assert interactive == configured_estimate + assert await catalog_task == default_estimate + async def test_unavailable_estimate_cache_expires(self) -> None: """A transient estimate failure is retried after the unavailable-result TTL.""" metadata = _make_scenario_metadata() @@ -376,7 +576,7 @@ async def test_unavailable_estimate_cache_expires(self) -> None: ) with ( - patch.object(ScenarioService, "__init__", lambda self: None), + patch.object(ScenarioService, "__init__", _initialize_test_service), patch("pyrit.backend.services.scenario_service._UNAVAILABLE_CACHE_TTL_SECONDS", 0), ): service = ScenarioService() @@ -402,7 +602,7 @@ async def test_estimate_cache_is_version_aware_and_bounded(self) -> None: scenario.get_default_run_size_estimate_async = AsyncMock(return_value=estimate) with ( - patch.object(ScenarioService, "__init__", lambda self: None), + patch.object(ScenarioService, "__init__", _initialize_test_service), patch("pyrit.backend.services.scenario_service._ESTIMATE_CACHE_SIZE", 1), ): service = ScenarioService() @@ -421,7 +621,7 @@ async def test_list_scenarios_preserves_disabled_baseline_policy(self) -> None: include_baseline_by_default=False, ) - with patch.object(ScenarioService, "__init__", lambda self: None): + with patch.object(ScenarioService, "__init__", _initialize_test_service): service = ScenarioService() service._registry = MagicMock() service._registry.get_all_registered_class_metadata.return_value = [metadata] @@ -437,7 +637,7 @@ async def test_list_scenarios_paginates_with_limit(self) -> None: _make_scenario_metadata(registry_name=f"test.scenario_{i}", class_name=f"Scenario{i}") for i in range(5) ] - with patch.object(ScenarioService, "__init__", lambda self: None): + with patch.object(ScenarioService, "__init__", _initialize_test_service): service = ScenarioService() service._registry = MagicMock() service._registry.get_all_registered_class_metadata.return_value = metadata_list @@ -459,7 +659,7 @@ async def test_list_scenarios_paginates_with_cursor(self) -> None: _make_scenario_metadata(registry_name=f"test.scenario_{i}", class_name=f"Scenario{i}") for i in range(5) ] - with patch.object(ScenarioService, "__init__", lambda self: None): + with patch.object(ScenarioService, "__init__", _initialize_test_service): service = ScenarioService() service._registry = MagicMock() service._registry.get_all_registered_class_metadata.return_value = metadata_list @@ -477,7 +677,7 @@ async def test_list_scenarios_last_page_has_more_false(self) -> None: _make_scenario_metadata(registry_name=f"test.scenario_{i}", class_name=f"Scenario{i}") for i in range(3) ] - with patch.object(ScenarioService, "__init__", lambda self: None): + with patch.object(ScenarioService, "__init__", _initialize_test_service): service = ScenarioService() service._registry = MagicMock() service._registry.get_all_registered_class_metadata.return_value = metadata_list @@ -506,7 +706,7 @@ async def test_configured_estimate_uses_shared_launch_resolution(self) -> None: objective_target = MagicMock() with ( - patch.object(ScenarioService, "__init__", lambda self: None), + patch.object(ScenarioService, "__init__", _initialize_test_service), patch.object( ScenarioConfigurationResolver, "resolve_target", return_value=objective_target ) as resolve_target, @@ -559,7 +759,7 @@ async def test_configured_estimate_rejects_incompatible_v4_jailbreak_technique(s introspection_instance._default_dataset_config = DatasetAttackConfiguration(dataset_names=["harmbench"]) scenario_class = MagicMock(return_value=introspection_instance) - with patch.object(ScenarioService, "__init__", lambda self: None): + with patch.object(ScenarioService, "__init__", _initialize_test_service): service = ScenarioService() service._registry = MagicMock() service._registry.get_registered_class_metadata.return_value = metadata @@ -586,7 +786,7 @@ async def test_configured_estimate_without_target_does_not_resolve_or_send_to_ta scenario_class = MagicMock(return_value=introspection_instance) with ( - patch.object(ScenarioService, "__init__", lambda self: None), + patch.object(ScenarioService, "__init__", _initialize_test_service), patch.object(ScenarioConfigurationResolver, "resolve_target") as resolve_target, ): service = ScenarioService() @@ -610,7 +810,7 @@ async def test_get_scenario_returns_matching_scenario(self) -> None: """Test that get returns the matching scenario.""" metadata = _make_scenario_metadata(registry_name="foundry.red_team_agent") - with patch.object(ScenarioService, "__init__", lambda self: None): + with patch.object(ScenarioService, "__init__", _initialize_test_service): service = ScenarioService() service._registry = MagicMock() service._registry.get_registered_class_metadata.return_value = metadata @@ -622,7 +822,7 @@ async def test_get_scenario_returns_matching_scenario(self) -> None: async def test_get_scenario_returns_none_for_missing(self) -> None: """Test that get returns None when scenario not found.""" - with patch.object(ScenarioService, "__init__", lambda self: None): + with patch.object(ScenarioService, "__init__", _initialize_test_service): service = ScenarioService() service._registry = MagicMock() service._registry.get_registered_class_metadata.return_value = None @@ -774,8 +974,11 @@ def test_get_scenario_returns_200(self, client: TestClient) -> None: aggregate_techniques=["all"], all_techniques=["role_play"], default_datasets=["airt_hate"], - default_run_size=ScenarioRunSizeEstimate( - estimated_attack_count=8, + default_run_size=ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Exact, + total_attack_count=8, + minimum_attack_count=8, + maximum_attack_count=8, components=[ ScenarioRunSizeComponent( label="Default technique sweep", @@ -796,7 +999,11 @@ def test_get_scenario_returns_200(self, client: TestClient) -> None: data = response.json() assert data["scenario_name"] == "foundry.red_team_agent" assert data["default_techniques"] == ["role_play"] - assert data["default_run_size"]["estimated_attack_count"] == 8 + assert data["default_run_size"]["version"] == 1 + assert data["default_run_size"]["status"] == "exact" + assert data["default_run_size"]["total_attack_count"] == 8 + assert data["default_run_size"]["minimum_attack_count"] == 8 + assert data["default_run_size"]["maximum_attack_count"] == 8 def test_get_scenario_returns_404_when_not_found(self, client: TestClient) -> None: """Test that GET /api/scenarios/catalog/{name} returns 404 when not found.""" @@ -836,7 +1043,9 @@ def test_estimate_scenario_returns_configured_projection(self, client: TestClien ) assert response.status_code == status.HTTP_200_OK - assert response.json()["estimated_attack_count"] == 12 + assert response.json()["total_attack_count"] == 12 + assert response.json()["minimum_attack_count"] is None + assert response.json()["maximum_attack_count"] is None request = mock_service.estimate_scenario_run_size_async.await_args.kwargs["request"] assert request.techniques == ["prompt_sending"] assert request.include_baseline is False @@ -957,7 +1166,7 @@ async def test_list_scenarios_includes_supported_parameters(self) -> None: ), ) - with patch.object(ScenarioService, "__init__", lambda self: None): + with patch.object(ScenarioService, "__init__", _initialize_test_service): service = ScenarioService() service._registry = MagicMock() service._registry.get_all_registered_class_metadata.return_value = [metadata] @@ -986,7 +1195,7 @@ async def test_scenario_with_no_parameters_has_empty_list(self) -> None: """Test that scenarios without parameters have empty supported_parameters.""" metadata = _make_scenario_metadata() - with patch.object(ScenarioService, "__init__", lambda self: None): + with patch.object(ScenarioService, "__init__", _initialize_test_service): service = ScenarioService() service._registry = MagicMock() service._registry.get_all_registered_class_metadata.return_value = [metadata] @@ -1016,7 +1225,7 @@ async def test_supported_parameters_with_none_default(self) -> None: ), ) - with patch.object(ScenarioService, "__init__", lambda self: None): + with patch.object(ScenarioService, "__init__", _initialize_test_service): service = ScenarioService() service._registry = MagicMock() service._registry.get_all_registered_class_metadata.return_value = [metadata] diff --git a/tests/unit/models/test_scenario_catalog.py b/tests/unit/models/test_scenario_catalog.py index 7e89b91ccf..9943beb7ff 100644 --- a/tests/unit/models/test_scenario_catalog.py +++ b/tests/unit/models/test_scenario_catalog.py @@ -7,62 +7,187 @@ from pydantic import ValidationError from pyrit.models import ( + ScenarioAdaptiveRunSizeDetails, ScenarioDatasetSizeCap, ScenarioDatasetSummary, ScenarioRunSizeComponent, ScenarioRunSizeEstimate, + ScenarioRunSizeEstimateCondition, ScenarioRunSizeEstimateRequest, ) +from pyrit.models.catalog import ( + ScenarioDefaultRunSizeEstimate, + ScenarioRunSizeEstimateStatus, + ScenarioRunSizeFactor, +) + + +def test_run_size_estimate_compatibility_alias_is_canonical_model() -> None: + """The initial DTO name remains an unambiguous alias of the versioned model.""" + assert ScenarioRunSizeEstimate is ScenarioDefaultRunSizeEstimate + + +def test_run_size_estimate_accepts_legacy_fields_and_serializes_canonically() -> None: + """Legacy constructors parse while the wire shape remains singular and versioned.""" + estimate = ScenarioRunSizeEstimate.model_validate( + { + "status": "exact", + "total": 2, + "components": [{"label": "Sweep", "count": 2}], + "datasets": [ + { + "name": "harmbench", + "seed_group_count": 100, + "selected_seed_group_count": 2, + } + ], + "caveat": "Legacy explanation.", + } + ) + + assert estimate.total == 2 + assert estimate.caveat == "Legacy explanation." + payload = estimate.model_dump(mode="json") + assert payload["version"] == 1 + assert payload["total_attack_count"] == 2 + assert payload["minimum_attack_count"] is None + assert payload["maximum_attack_count"] is None + assert payload["note"] == "Legacy explanation." + assert payload["datasets"][0]["logical_seed_group_count"] == 100 + assert "total" not in payload + assert "caveat" not in payload + assert "seed_group_count" not in payload["datasets"][0] + +def test_run_size_estimate_normalizes_legacy_componentless_exact_total() -> None: + estimate = ScenarioRunSizeEstimate.model_validate({"status": "exact", "total": 2}) -def test_run_size_estimate_requires_available_total_to_match_components() -> None: - """Available estimates require an additive component total.""" + assert estimate.total_attack_count == 2 + assert estimate.components == [ + ScenarioRunSizeComponent( + label="Legacy total", + count=2, + note="Normalized from a legacy component-less estimate.", + ) + ] + + +def test_exact_default_run_size_requires_component_total() -> None: + """Exact estimates reject totals that do not match their additive components.""" with pytest.raises(ValidationError, match="components total 6, not 7"): - ScenarioRunSizeEstimate( - estimated_attack_count=7, + ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Exact, + total_attack_count=7, + components=[ + ScenarioRunSizeComponent( + label="Techniques", + count=6, + factors=[ + ScenarioRunSizeFactor(label="seed groups", count=3), + ScenarioRunSizeFactor(label="techniques", count=2), + ], + ) + ], + ) + + +@pytest.mark.parametrize("field_name", ["minimum_attack_count", "maximum_attack_count"]) +def test_exact_default_run_size_requires_bounds_to_match_total(field_name: str) -> None: + """Exact estimates reject bounds that disagree with their authoritative total.""" + with pytest.raises(ValidationError, match=f"{field_name} to equal total_attack_count"): + ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Exact, + total_attack_count=6, components=[ScenarioRunSizeComponent(label="Techniques", count=6)], + **{field_name: 5}, ) -def test_run_size_estimate_allows_unavailable_count_with_components() -> None: - """Unavailable estimates retain useful candidate components and an explanatory note.""" - estimate = ScenarioRunSizeEstimate( - components=[ScenarioRunSizeComponent(label="Candidate techniques", count=6)], - note="The final count depends on target capabilities.", - ) +def test_default_run_size_requires_ordered_nonnegative_bounds() -> None: + """Conditional estimate bounds remain nonnegative and ordered.""" + with pytest.raises(ValidationError, match="greater than or equal to 0"): + ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Conditional, + minimum_attack_count=-1, + ) + + with pytest.raises(ValidationError, match="minimum_attack_count must be less than or equal"): + ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Conditional, + minimum_attack_count=20, + maximum_attack_count=12, + ) + - assert estimate.estimated_attack_count is None - assert estimate.components[0].count == 6 +def test_conditional_default_run_size_allows_unknown_bounds() -> None: + """Conditional estimates may remain unbounded when no truthful range is available.""" + estimate = ScenarioDefaultRunSizeEstimate(status=ScenarioRunSizeEstimateStatus.Conditional) + assert estimate.minimum_attack_count is None + assert estimate.maximum_attack_count is None -def test_run_size_estimate_serializes_canonical_api_shape() -> None: - """The estimate exposes only the available count and additive components.""" - estimate = ScenarioRunSizeEstimate( - estimated_attack_count=6, - components=[ScenarioRunSizeComponent(label="Techniques", count=6)], + +def test_run_size_component_requires_factor_product() -> None: + """Components reject counts that disagree with their ordered formula factors.""" + with pytest.raises(ValidationError, match="factor product \\(6\\)"): + ScenarioRunSizeComponent( + label="Techniques", + count=7, + factors=[ + ScenarioRunSizeFactor(label="seed groups", count=3), + ScenarioRunSizeFactor(label="techniques", count=2), + ], + ) + + +def test_default_run_size_serializes_versioned_api_shape() -> None: + """The estimate exposes stable status, total, component, and factor fields.""" + estimate = ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Exact, + total_attack_count=6, + components=[ + ScenarioRunSizeComponent( + label="Techniques", + count=6, + factors=[ + ScenarioRunSizeFactor(label="seed groups", count=3), + ScenarioRunSizeFactor(label="techniques", count=2), + ], + ) + ], ) assert estimate.model_dump(mode="json") == { - "estimated_attack_count": 6, + "version": 1, + "status": "exact", + "total_attack_count": 6, "minimum_attack_count": None, "maximum_attack_count": None, + "condition": None, "components": [ { "label": "Techniques", "count": 6, + "factors": [ + {"label": "seed groups", "count": 3}, + {"label": "techniques", "count": 2}, + ], "note": None, "is_baseline": False, + "condition": None, } ], "datasets": [], "effective_parameters": {}, + "adaptive_details": None, "note": None, + "retries_included": False, } def test_run_size_estimate_rejects_inverted_bounds() -> None: """The minimum estimate cannot exceed the maximum estimate.""" - with pytest.raises(ValidationError, match="Minimum attack count cannot exceed maximum attack count"): + with pytest.raises(ValidationError, match="minimum_attack_count must be less than or equal"): ScenarioRunSizeEstimate(minimum_attack_count=8, maximum_attack_count=4) @@ -70,13 +195,81 @@ def test_unavailable_run_size_estimate_has_no_count() -> None: """The unavailable factory communicates that a count cannot be calculated.""" estimate = ScenarioRunSizeEstimate.unavailable() - assert estimate.estimated_attack_count is None - assert estimate.note == "Default-run size estimate is unavailable." + assert estimate.status is ScenarioRunSizeEstimateStatus.Unavailable + assert estimate.total_attack_count is None + + +def test_adaptive_run_size_details_serialize_derived_attempt_bounds() -> None: + """Adaptive estimates expose progress objectives and underlying attempt bounds separately.""" + details = ScenarioAdaptiveRunSizeDetails( + objective_count=21, + selected_candidate_technique_count=2, + candidate_technique_count=2, + max_attempts_per_objective=3, + techniques_per_objective_upper_bound=2, + technique_attempt_count_upper_bound=42, + ) + + assert details.model_dump(mode="json") == { + "objective_count": 21, + "selected_candidate_technique_count": 2, + "candidate_technique_count": 2, + "max_attempts_per_objective": 3, + "techniques_per_objective_upper_bound": 2, + "technique_attempt_count_upper_bound": 42, + "stop_on_first_success": True, + "compatibility_may_reduce_attempts": True, + } + + +def test_adaptive_run_size_details_reject_inconsistent_attempt_bounds() -> None: + """Adaptive work bounds cannot drift from the selected pool and configured cap.""" + with pytest.raises(ValidationError, match="min\\(candidate_technique_count, max_attempts_per_objective\\)"): + ScenarioAdaptiveRunSizeDetails( + objective_count=21, + selected_candidate_technique_count=2, + candidate_technique_count=2, + max_attempts_per_objective=3, + techniques_per_objective_upper_bound=3, + technique_attempt_count_upper_bound=63, + ) + + +def test_adaptive_run_size_details_accept_legacy_version_one_payload() -> None: + """Version-one payloads without the additive selected count remain readable.""" + details = ScenarioAdaptiveRunSizeDetails.model_validate( + { + "objective_count": 21, + "candidate_technique_count": 2, + "max_attempts_per_objective": 3, + "techniques_per_objective_upper_bound": 2, + "technique_attempt_count_upper_bound": 42, + } + ) + + assert details.selected_candidate_technique_count == 2 + + +def test_adaptive_run_size_details_reject_more_compatible_than_selected_candidates() -> None: + """Resolved compatible candidates cannot exceed the concrete selected pool.""" + with pytest.raises(ValidationError, match="cannot exceed selected_candidate_technique_count"): + ScenarioAdaptiveRunSizeDetails( + objective_count=21, + selected_candidate_technique_count=2, + candidate_technique_count=3, + max_attempts_per_objective=3, + techniques_per_objective_upper_bound=3, + technique_attempt_count_upper_bound=63, + ) -def test_estimate_exposes_dataset_counts_structurally() -> None: - """Effective dataset selection remains machine-readable.""" - estimate = ScenarioRunSizeEstimate( +def test_conditional_estimate_exposes_dataset_counts_structurally() -> None: + """Conditionality and effective dataset selection are machine-readable.""" + estimate = ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Conditional, + minimum_attack_count=12, + maximum_attack_count=20, + condition=ScenarioRunSizeEstimateCondition.TargetCapabilities, datasets=[ ScenarioDatasetSummary( name="harmbench", @@ -93,11 +286,16 @@ def test_estimate_exposes_dataset_counts_structurally() -> None: ], ) ], - note="The final count depends on target capabilities.", + note="The final total depends on target capabilities.", ) - assert estimate.estimated_attack_count is None - assert estimate.model_dump(mode="json")["datasets"] == [ + payload = estimate.model_dump(mode="json") + assert payload["status"] == "conditional" + assert payload["total_attack_count"] is None + assert payload["minimum_attack_count"] == 12 + assert payload["maximum_attack_count"] == 20 + assert payload["condition"] == "target_capabilities" + assert payload["datasets"] == [ { "name": "harmbench", "kind": "dataset", @@ -114,6 +312,7 @@ def test_estimate_exposes_dataset_counts_structurally() -> None: ], } ] + assert payload["retries_included"] is False def test_estimate_request_reuses_dataset_filter_validation() -> None: diff --git a/tests/unit/registry/test_scenario_registry.py b/tests/unit/registry/test_scenario_registry.py index 649ef9beb7..da812520fe 100644 --- a/tests/unit/registry/test_scenario_registry.py +++ b/tests/unit/registry/test_scenario_registry.py @@ -3,12 +3,21 @@ """Tests for ScenarioRegistry._build_metadata and create_and_initialize_async.""" -from unittest.mock import AsyncMock, MagicMock +from typing import Literal +from unittest.mock import AsyncMock, MagicMock, patch import pytest -from pyrit.registry.components.scenario_registry import ScenarioRegistry -from pyrit.scenario.core import BaselineAttackPolicy, ScenarioTechnique +from pyrit.registry import ScenarioRegistry +from pyrit.scenario import ( + BaselineAttackPolicy, + CompoundDatasetAttackConfiguration, + DatasetAttackConfiguration, + ScenarioTechnique, +) +from pyrit.scenario.scenarios.adaptive import TextAdaptive +from pyrit.scenario.scenarios.airt import Psychosocial +from pyrit.scenario.scenarios.garak import WebInjection class _NotNoArgScenario: @@ -61,6 +70,10 @@ def _resolve_scenario_techniques(self, *, scenario_techniques): """Resolve the concrete defaults.""" return _MetadataTechnique.resolve(scenario_techniques, default=self._default_technique) + def get_dataset_size_limit_override_scope(self) -> Literal["per_dataset"]: + """Return the test scenario's conventional single-dataset override scope.""" + return "per_dataset" + class _MarkdownMetadataScenario(_MetadataScenario): """ @@ -104,6 +117,92 @@ def test_build_metadata_expands_ordered_default_techniques() -> None: } +@pytest.mark.parametrize( + ("configuration", "declared_override_scope", "default_scope", "default_count", "override_scope"), + [ + (DatasetAttackConfiguration(dataset_names=["sample"]), "per_dataset", "none", None, "per_dataset"), + ( + DatasetAttackConfiguration(dataset_names=["one", "two"]), + "per_dataset", + "none", + None, + "per_dataset", + ), + ( + DatasetAttackConfiguration(dataset_names=["one", "two"]), + "unsupported", + "none", + None, + "unsupported", + ), + ( + DatasetAttackConfiguration(dataset_names=["one", "two"], max_dataset_size=6), + "combined", + "combined", + 6, + "combined", + ), + ( + CompoundDatasetAttackConfiguration.per_dataset( + dataset_names=["one", "two"], + max_dataset_size=4, + ), + "per_dataset", + "per_dataset", + 4, + "per_dataset", + ), + ( + CompoundDatasetAttackConfiguration( + configurations=[ + DatasetAttackConfiguration(dataset_names=["one"], max_dataset_size=3), + DatasetAttackConfiguration(dataset_names=["two"], max_dataset_size=4), + ] + ), + "per_dataset", + "heterogeneous", + None, + "per_dataset", + ), + ], +) +def test_build_dataset_size_limit_normalizes_configuration_semantics( + configuration: DatasetAttackConfiguration, + declared_override_scope: Literal["per_dataset", "combined", "unsupported"], + default_scope: str, + default_count: int | None, + override_scope: str, +) -> None: + """Catalog limit metadata preserves no-cap, combined, per-dataset, and heterogeneous defaults.""" + limit = ScenarioRegistry._build_dataset_size_limit( + default_dataset_config=configuration, + override_scope=declared_override_scope, + ) + + assert limit.default_scope == default_scope + assert limit.default_count == default_count + assert limit.override_scope == override_scope + + +def test_specialized_scenarios_declare_nonstandard_dataset_override_semantics() -> None: + """Catalog metadata can remain truthful when a scenario reshapes or ignores generic dataset caps.""" + assert Psychosocial.DATASET_SIZE_LIMIT_OVERRIDE_SCOPE == "per_dataset" + assert WebInjection.DATASET_SIZE_LIMIT_OVERRIDE_SCOPE == "unsupported" + + +def test_text_adaptive_metadata_exposes_per_dataset_default_limit() -> None: + """TextAdaptive publishes its canonical four-objective child cap without scenario-name special cases.""" + with ( + patch.object(TextAdaptive, "_get_default_objective_scorer", return_value=MagicMock()), + patch("pyrit.scenario.core.scenario.CentralMemory.get_memory_instance", return_value=MagicMock()), + ): + metadata = ScenarioRegistry()._build_metadata("adaptive.text_adaptive", TextAdaptive) + + assert metadata.dataset_size_limit.default_scope == "per_dataset" + assert metadata.dataset_size_limit.default_count == 4 + assert metadata.dataset_size_limit.override_scope == "per_dataset" + + def test_build_metadata_preserves_structured_markdown_separately() -> None: """Scenario metadata keeps plain compatibility text and Markdown source.""" metadata = ScenarioRegistry()._build_metadata("markdown", _MarkdownMetadataScenario) diff --git a/tests/unit/scenario/airt/test_jailbreak.py b/tests/unit/scenario/airt/test_jailbreak.py index 284736b104..580a7c226b 100644 --- a/tests/unit/scenario/airt/test_jailbreak.py +++ b/tests/unit/scenario/airt/test_jailbreak.py @@ -3,6 +3,7 @@ """Tests for the Jailbreak class.""" +import logging from typing import Any from unittest.mock import AsyncMock, MagicMock, patch @@ -18,10 +19,10 @@ SeedObjective, SeedPrompt, ) +from pyrit.models.catalog import ScenarioRunSizeEstimateStatus from pyrit.prompt_target import PromptTarget -from pyrit.registry import TargetRegistry +from pyrit.registry import ScenarioRegistry, TargetRegistry from pyrit.registry.components.attack_technique_registry import AttackTechniqueRegistry -from pyrit.registry.components.scenario_registry import ScenarioRegistry from pyrit.scenario.core import BaselineAttackPolicy from pyrit.scenario.core.attack_technique_factory import AttackTechniqueFactory from pyrit.scenario.scenarios.airt.jailbreak import ( @@ -251,9 +252,11 @@ async def test_run_size_is_conditional_when_system_delivery_target_is_not_select ) estimate = await scenario.get_run_size_estimate_async(target_is_configured=False) - assert estimate.estimated_attack_count is None + assert estimate.status is ScenarioRunSizeEstimateStatus.Conditional + assert estimate.total_attack_count is None assert estimate.minimum_attack_count == 2 assert estimate.maximum_attack_count == 4 + assert estimate.condition.value == "target_capabilities" assert [component.label for component in estimate.components] == [ "Inline jailbreak delivery", "Native system-prompt jailbreak delivery", diff --git a/tests/unit/scenario/core/test_dataset_configuration.py b/tests/unit/scenario/core/test_dataset_configuration.py index 3c3ad9c25c..5e1aabd103 100644 --- a/tests/unit/scenario/core/test_dataset_configuration.py +++ b/tests/unit/scenario/core/test_dataset_configuration.py @@ -262,6 +262,16 @@ async def test_max_sample_is_a_single_global_budget(self, mock_memory: MagicMock result = await config.get_attack_groups_by_dataset_async() assert sum(len(groups) for groups in result.values()) == 2 + async def test_estimate_resolution_fetches_full_and_sampled_groups_once(self, mock_memory: MagicMock) -> None: + mock_memory.get_seeds.return_value = make_objectives("a", "b", "c") + config = DatasetAttackConfiguration(dataset_names=["d1"], max_dataset_size=1) + + full, selected = await config.resolve_attack_groups_for_estimate_async() + + assert len(full["d1"]) == 3 + assert len(selected["d1"]) == 1 + mock_memory.get_seeds.assert_called_once() + async def test_loud_raise_when_a_dataset_is_empty(self, mock_memory: MagicMock) -> None: mock_memory.get_seeds.side_effect = [make_objectives("a"), []] config = DatasetAttackConfiguration(dataset_names=["d1", "d2"], auto_fetch=False) @@ -507,6 +517,95 @@ def test_per_dataset_builds_one_child_per_name(self) -> None: assert [child.dataset_names for child in config._configurations] == [["d1"], ["d2"]] assert all(child.max_dataset_size == 4 for child in config._configurations) + def test_with_dataset_names_preserves_per_dataset_defaults(self) -> None: + config = CompoundDatasetAttackConfiguration.per_dataset( + dataset_names=["d1", "d2"], + max_dataset_size=4, + filters={"harm_categories": ["original"]}, + ) + + overridden = config.with_dataset_names( + dataset_names=["selected"], + filters={"data_types": ["text"]}, + ) + + assert overridden.dataset_names == ["selected"] + assert len(overridden._configurations) == 1 + assert overridden._configurations[0].max_dataset_size == 4 + assert overridden._configurations[0].filters == { + "harm_categories": ["original"], + "data_types": ["text"], + } + assert overridden._configurations[0]._validators == config._configurations[0]._validators + assert overridden._validators == config._validators + assert config.dataset_names == ["d1", "d2"] + + def test_with_dataset_names_rejects_shaped_children(self) -> None: + class _ShapedDatasetConfiguration(DatasetAttackConfiguration): + pass + + config = CompoundDatasetAttackConfiguration( + configurations=[_ShapedDatasetConfiguration(dataset_names=["d1"])], + ) + + with pytest.raises(TypeError, match="homogeneous single-dataset"): + config.with_dataset_names(dataset_names=["selected"]) + + def test_with_dataset_names_rejects_duplicates(self) -> None: + config = CompoundDatasetAttackConfiguration.per_dataset(dataset_names=["d1", "d2"]) + + with pytest.raises(ValueError, match="cannot contain duplicates"): + config.with_dataset_names(dataset_names=["selected", "selected"]) + + def test_with_dataset_names_rejects_different_child_validators(self) -> None: + first_validator = require_min_size(1) + second_validator = require_min_size(2) + config = CompoundDatasetAttackConfiguration( + configurations=[ + DatasetAttackConfiguration(dataset_names=["d1"], validators=[first_validator]), + DatasetAttackConfiguration(dataset_names=["d2"], validators=[second_validator]), + ], + ) + + with pytest.raises(TypeError, match="shared caps, filters, validators"): + config.with_dataset_names(dataset_names=["selected"]) + + def test_with_dataset_names_supports_shared_unhashable_validator(self) -> None: + class _UnhashableValidator: + __hash__ = None + + def __call__(self, resolved: ResolvedDataset) -> None: + del resolved + + validator = _UnhashableValidator() + config = CompoundDatasetAttackConfiguration.per_dataset( + dataset_names=["d1", "d2"], + validators=[validator], + ) + + overridden = config.with_dataset_names(dataset_names=["selected"]) + + assert overridden._configurations[0]._custom_validators == [validator] + + def test_update_child_max_dataset_size_preserves_shaped_children(self) -> None: + class _ShapedDatasetConfiguration(DatasetAttackConfiguration): + pass + + config = CompoundDatasetAttackConfiguration( + configurations=[ + _ShapedDatasetConfiguration(dataset_names=["d1"], max_dataset_size=4), + _ShapedDatasetConfiguration(dataset_names=["d2"], max_dataset_size=4), + ], + ) + + config.update_child_max_dataset_size(max_dataset_size=2) + + assert [type(child) for child in config._configurations] == [ + _ShapedDatasetConfiguration, + _ShapedDatasetConfiguration, + ] + assert [child.max_dataset_size for child in config._configurations] == [2, 2] + def test_size_caps_report_child_and_combined_limits(self) -> None: """Planning metadata explains independent child caps and the final compound cap.""" config = CompoundDatasetAttackConfiguration.per_dataset(dataset_names=["d1", "d2"], max_dataset_size=4) diff --git a/tests/unit/scenario/scenarios/adaptive/test_dispatcher.py b/tests/unit/scenario/scenarios/adaptive/test_dispatcher.py index eceb2b5cf9..b33a769829 100644 --- a/tests/unit/scenario/scenarios/adaptive/test_dispatcher.py +++ b/tests/unit/scenario/scenarios/adaptive/test_dispatcher.py @@ -18,6 +18,8 @@ ) from pyrit.scenario.scenarios.adaptive.dispatcher import ( ADAPTIVE_ATTEMPT_LABEL, + ADAPTIVE_TECHNIQUE_ID_LABEL, + ADAPTIVE_TECHNIQUE_NAME_LABEL, AdaptiveTechniqueDispatcher, TechniqueBundle, ) @@ -127,6 +129,10 @@ async def test_builds_sequential_attack(self, target, seed_group): # 1-based per-attempt label stamped on each child assert attack._child_attacks[0].memory_labels[ADAPTIVE_ATTEMPT_LABEL] == "1" assert attack._child_attacks[1].memory_labels[ADAPTIVE_ATTEMPT_LABEL] == "2" + assert attack._child_attacks[0].memory_labels[ADAPTIVE_TECHNIQUE_ID_LABEL] == "a" + assert attack._child_attacks[1].memory_labels[ADAPTIVE_TECHNIQUE_ID_LABEL] == "b" + assert attack._child_attacks[0].memory_labels[ADAPTIVE_TECHNIQUE_NAME_LABEL] == "a" + assert attack._child_attacks[1].memory_labels[ADAPTIVE_TECHNIQUE_NAME_LABEL] == "b" # default policy is FIRST_SUCCESS assert attack._completion_policy is SequenceCompletionPolicy.FIRST_SUCCESS @@ -239,37 +245,25 @@ async def test_merges_real_system_prompt_technique_onto_user_turn_at_sequence_ze @pytest.mark.usefixtures("patch_central_database") -class TestEvalHashRoundTrip: +class TestRegisteredTechniqueIdentityRoundTrip: """ - Pin the load-bearing invariant that ``compute_inner_attack_eval_hash`` - (used by ``AdaptiveScenario._build_techniques_dict`` to key the - ``techniques`` dict and by the selector to look up historical stats) - equals the ``eval_hash`` the executor stamps on persisted child rows. - - If the prediction helper and the write path ever drift (e.g. a new - field is added to the eval-hash rule on one side only), the selector - silently reads zero history for every technique and epsilon-greedy - degrades to random with no error. This test runs a real - ``PromptSendingAttack`` through the dispatcher's ``SequentialAttack`` - end-to-end and asserts the round-trip holds. + Pin the selector identity labels through real child-result persistence. """ - async def test_predicted_hash_matches_persisted_row(self, sqlite_instance): + async def test_registered_identity_and_name_are_persisted(self, sqlite_instance): from pyrit.executor.attack.single_turn.prompt_sending import PromptSendingAttack from pyrit.memory.memory_models import AttackResultEntry from pyrit.models import AttackSeedGroup, SeedObjective - from pyrit.models.identifiers import compute_inner_attack_eval_hash from tests.unit.mocks import MockPromptTarget live_target = MockPromptTarget() attack = PromptSendingAttack(objective_target=live_target) - predicted_hash = compute_inner_attack_eval_hash(attack=attack) - - bundles = {predicted_hash: TechniqueBundle(attack=attack, name="prompt_sending")} + technique_identifier = "factory-identity-hash" + bundles = {technique_identifier: TechniqueBundle(attack=attack, name="prompt_sending")} dispatcher = AdaptiveTechniqueDispatcher( objective_target=live_target, techniques=bundles, - selector=_StubSelector(technique_order=[predicted_hash]), + selector=_StubSelector(technique_order=[technique_identifier]), max_attempts_per_objective=1, ) @@ -280,8 +274,6 @@ async def test_predicted_hash_matches_persisted_row(self, sqlite_instance): with sqlite_instance.get_session() as session: rows = session.query(AttackResultEntry).all() - # Drill into the persisted envelope to find rows whose inner attack is PromptSendingAttack, - # then assert the eval_hash on those rows matches what the selector predicted. matching_rows = [ r for r in rows @@ -297,10 +289,5 @@ async def test_predicted_hash_matches_persisted_row(self, sqlite_instance): f"Expected at least one persisted row whose inner attack is PromptSendingAttack; " f"found rows: {[(r.id, r.atomic_attack_identifier) for r in rows]}" ) - for row in matching_rows: - stamped_hash = row.atomic_attack_identifier["eval_hash"] - assert stamped_hash == predicted_hash, ( - f"Selector-side eval_hash ({predicted_hash}) drifted from executor-stamped " - f"eval_hash ({stamped_hash}) on persisted row {row.id}. " - f"compute_inner_attack_eval_hash and AtomicAttackIdentifier.build must agree." - ) + assert all(row.labels[ADAPTIVE_TECHNIQUE_ID_LABEL] == technique_identifier for row in matching_rows) + assert all(row.labels[ADAPTIVE_TECHNIQUE_NAME_LABEL] == "prompt_sending" for row in matching_rows) diff --git a/tests/unit/scenario/scenarios/adaptive/test_epsilon_greedy.py b/tests/unit/scenario/scenarios/adaptive/test_epsilon_greedy.py index 21144721f5..23d4ee12f9 100644 --- a/tests/unit/scenario/scenarios/adaptive/test_epsilon_greedy.py +++ b/tests/unit/scenario/scenarios/adaptive/test_epsilon_greedy.py @@ -6,14 +6,15 @@ import pytest from pyrit.analytics.result_analysis import AttackStats -from pyrit.scenario.scenarios.adaptive.selectors import ( +from pyrit.scenario.scenarios.adaptive import ( + AdaptiveTechniqueIdentifier, EpsilonGreedyTechniqueSelector, SelectorScope, ) TECHNIQUES = ["a", "b", "c", "d"] -_COMPUTE_PATH = "pyrit.scenario.scenarios.adaptive.selectors.epsilon_greedy.compute_technique_stats" +_COMPUTE_PATH = "pyrit.scenario.scenarios.adaptive.selectors.epsilon_greedy.compute_labeled_technique_stats" def _seeded_selector(*, epsilon: float = 0.0, random_seed: int = 0) -> EpsilonGreedyTechniqueSelector: @@ -66,7 +67,7 @@ def test_init_rejects_out_of_range_epsilon(self, bad_epsilon): class TestEpsilonGreedyTechniqueSelectorSelect: @patch( - "pyrit.scenario.scenarios.adaptive.selectors.epsilon_greedy.compute_technique_stats", + "pyrit.scenario.scenarios.adaptive.selectors.epsilon_greedy.compute_labeled_technique_stats", side_effect=_empty_rates, ) async def test_select_empty_techniques_raises(self, _mock): @@ -75,7 +76,7 @@ async def test_select_empty_techniques_raises(self, _mock): await selector.select_async(technique_identifiers=[], objective="obj") @patch( - "pyrit.scenario.scenarios.adaptive.selectors.epsilon_greedy.compute_technique_stats", + "pyrit.scenario.scenarios.adaptive.selectors.epsilon_greedy.compute_labeled_technique_stats", side_effect=_empty_rates, ) async def test_select_all_unseen_ties_resolved_randomly(self, _mock): @@ -88,7 +89,7 @@ async def test_select_all_unseen_ties_resolved_randomly(self, _mock): assert winners.issubset(set(TECHNIQUES)) @patch( - "pyrit.scenario.scenarios.adaptive.selectors.epsilon_greedy.compute_technique_stats", + "pyrit.scenario.scenarios.adaptive.selectors.epsilon_greedy.compute_labeled_technique_stats", side_effect=_rates_with_winner("b"), ) async def test_select_exploits_clear_winner(self, _mock): @@ -98,7 +99,7 @@ async def test_select_exploits_clear_winner(self, _mock): assert result[0] == "b" @patch( - "pyrit.scenario.scenarios.adaptive.selectors.epsilon_greedy.compute_technique_stats", + "pyrit.scenario.scenarios.adaptive.selectors.epsilon_greedy.compute_labeled_technique_stats", side_effect=_empty_rates, ) async def test_select_epsilon_one_is_pure_random(self, _mock): @@ -110,7 +111,7 @@ async def test_select_epsilon_one_is_pure_random(self, _mock): assert picks == set(TECHNIQUES) @patch( - "pyrit.scenario.scenarios.adaptive.selectors.epsilon_greedy.compute_technique_stats", + "pyrit.scenario.scenarios.adaptive.selectors.epsilon_greedy.compute_labeled_technique_stats", side_effect=_empty_rates, ) async def test_select_returns_multiple_techniques(self, _mock): @@ -120,7 +121,7 @@ async def test_select_returns_multiple_techniques(self, _mock): assert len(set(result)) == 3 # no duplicates @patch( - "pyrit.scenario.scenarios.adaptive.selectors.epsilon_greedy.compute_technique_stats", + "pyrit.scenario.scenarios.adaptive.selectors.epsilon_greedy.compute_labeled_technique_stats", side_effect=_empty_rates, ) async def test_select_caps_at_available_techniques(self, _mock): @@ -130,6 +131,18 @@ async def test_select_caps_at_available_techniques(self, _mock): class TestEpsilonGreedySelectorScope: + @patch(_COMPUTE_PATH, side_effect=_empty_rates) + async def test_forwards_full_technique_eval_hash_for_cross_scenario_history(self, mock_compute): + arm = AdaptiveTechniqueIdentifier( + factory_hash="factory-hash", + technique_eval_hash="full-technique-eval-hash", + ).serialize() + selector = _seeded_selector() + + await selector.select_async(technique_identifiers=[arm], objective="obj") + + assert mock_compute.call_args.kwargs["technique_eval_hashes_by_identifier"] == {arm: "full-technique-eval-hash"} + @patch(_COMPUTE_PATH, side_effect=_empty_rates) async def test_default_scope_passes_none_scenario_result_id(self, mock_compute): selector = _seeded_selector() diff --git a/tests/unit/scenario/scenarios/adaptive/test_technique_identity.py b/tests/unit/scenario/scenarios/adaptive/test_technique_identity.py new file mode 100644 index 0000000000..c491d6a1f7 --- /dev/null +++ b/tests/unit/scenario/scenarios/adaptive/test_technique_identity.py @@ -0,0 +1,23 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from pyrit.scenario.scenarios.adaptive.technique_identity import ( + AdaptiveTechniqueIdentifier, + get_history_eval_hash, +) + + +def test_adaptive_technique_identifier_round_trip() -> None: + identifier = AdaptiveTechniqueIdentifier( + factory_hash="factory-hash", + technique_eval_hash="technique-eval-hash", + ) + + serialized = identifier.serialize() + + assert AdaptiveTechniqueIdentifier.parse(serialized) == identifier + assert get_history_eval_hash(technique_identifier=serialized) == "technique-eval-hash" + + +def test_history_eval_hash_falls_back_for_custom_selector_identifier() -> None: + assert get_history_eval_hash(technique_identifier="custom-arm") == "custom-arm" diff --git a/tests/unit/scenario/scenarios/adaptive/test_text_adaptive.py b/tests/unit/scenario/scenarios/adaptive/test_text_adaptive.py index 93a9929b20..0483b6a30a 100644 --- a/tests/unit/scenario/scenarios/adaptive/test_text_adaptive.py +++ b/tests/unit/scenario/scenarios/adaptive/test_text_adaptive.py @@ -11,17 +11,29 @@ import pytest -from pyrit.models import AttackSeedGroup, SeedObjective +from pyrit.models import AttackSeedGroup, ScenarioDatasetSummary, SeedObjective from pyrit.models.identifiers import ComponentIdentifier from pyrit.prompt_target import PromptTarget from pyrit.registry.components.attack_technique_registry import AttackTechniqueRegistry from pyrit.scenario.core.dataset_configuration import CompoundDatasetAttackConfiguration from pyrit.scenario.core.scenario import BaselineAttackPolicy from pyrit.scenario.scenarios.adaptive.dispatcher import AdaptiveTechniqueDispatcher +from pyrit.scenario.scenarios.adaptive.technique_identity import AdaptiveTechniqueIdentifier from pyrit.scenario.scenarios.adaptive.text_adaptive import TextAdaptive from pyrit.score import TrueFalseScorer _MOCK_MANY_SHOT_EXAMPLES = [{"question": f"q{i}", "answer": f"a{i}"} for i in range(100)] +_LIGHT_TECHNIQUES = { + "role_play_movie_script", + "role_play_video_game", + "role_play_trivia_game", + "role_play_persuasion", + "role_play_persuasion_written", + "many_shot", + "red_teaming", + "context_compliance", + "flip", +} def _mock_id(name: str) -> ComponentIdentifier: @@ -82,30 +94,43 @@ def _make_seed_group(*, value: str, harm_categories: list[str] | None = None) -> return AttackSeedGroup(seeds=[SeedObjective(value=value, harm_categories=harm_categories)]) -def _make_fake_factory(*, seed_technique=None, adversarial_chat=None, scoring_config_type=None) -> MagicMock: +def _make_fake_factory( + *, + seed_technique=None, + adversarial_chat=None, + scoring_config_type=None, + attack_identifier: ComponentIdentifier | None = None, + factory_identifier: ComponentIdentifier | None = None, + technique_identifier: ComponentIdentifier | None = None, +) -> MagicMock: """Return a stub attack-technique factory that produces a fake ``AttackTechnique``. Mocks the surface ``AdaptiveScenario._build_techniques_dict`` consumes - (``factory.create(...)``, ``factory.adversarial_chat``, and - ``factory.scoring_config_type``). Each call assigns a unique fake - attack identifier (via a fresh UUID) so the bundle dict keys (eval - hashes) don't collide across calls — no shared mutable test state, so - test execution order doesn't shift hash values. + (``factory.create(...)``, ``factory.get_identifier()``, + ``factory.adversarial_chat``, and ``factory.scoring_config_type``). + Each call assigns unique attack and factory identities unless a test + deliberately supplies shared identities. """ fake_id = uuid.uuid4().hex[:8] fake_technique = MagicMock() fake_attack = MagicMock(name=f"fake-attack-technique-{fake_id}") - fake_attack.get_identifier.return_value = ComponentIdentifier( - class_name=f"FakeAttack{fake_id}", - class_module="test_text_adaptive", + fake_attack.get_identifier.return_value = attack_identifier or ComponentIdentifier( + class_name=f"FakeAttack{fake_id}", class_module="test_text_adaptive" ) fake_technique.attack = fake_attack fake_technique.seed_technique = seed_technique + fake_technique.get_identifier.return_value = technique_identifier or ComponentIdentifier( + class_name=f"FakeTechnique{fake_id}", class_module="test_text_adaptive" + ) factory = MagicMock() factory.create.return_value = fake_technique factory.adversarial_chat = adversarial_chat + factory.uses_adversarial = adversarial_chat is not None factory.scoring_config_type = scoring_config_type + factory.get_identifier.return_value = factory_identifier or ComponentIdentifier( + class_name=f"FakeFactory{fake_id}", class_module="test_text_adaptive" + ) return factory @@ -190,6 +215,17 @@ async def test_binding_dataset_cap_exposes_attack_range( assert estimate.minimum_attack_count == 2 assert estimate.maximum_attack_count == 4 + def test_max_attempts_parameter_distinguishes_techniques_from_retries(self): + parameter = next( + parameter + for parameter in TextAdaptive.additional_parameters() + if parameter.name == "max_attempts_per_objective" + ) + assert parameter.default == 3 + assert "different compatible techniques" in parameter.description + assert "stopping after the first success" in parameter.description + assert "separate from retries" in parameter.description + def test_get_technique_class_is_cached(self): cls_a = TextAdaptive.get_technique_class() cls_b = TextAdaptive.get_technique_class() @@ -394,6 +430,139 @@ async def test_techniques_with_seed_technique_are_kept(self, mock_objective_targ assert "role_play_movie_script" in technique_names assert "many_shot" in technique_names + @pytest.mark.parametrize(("max_attempts", "expected_attempts"), [(4, 84), (5, 105)]) + async def test_light_keeps_nine_distinct_factory_arms_and_attempt_bound( + self, + mock_objective_target, + mock_objective_scorer, + max_attempts, + expected_attempts, + ): + shared_attack_identifier = _mock_id("SharedPromptSendingAttack") + factories = { + name: _make_fake_factory( + attack_identifier=shared_attack_identifier, + factory_identifier=_mock_id(f"Factory_{name}"), + ) + for name in _LIGHT_TECHNIQUES + } + groups = {"adaptive": [_make_seed_group(value=f"obj-{index}") for index in range(21)]} + summaries = [ + ScenarioDatasetSummary( + name="adaptive", + logical_seed_group_count=21, + selected_seed_group_count=21, + ) + ] + scenario = TextAdaptive(objective_scorer=mock_objective_scorer) + technique_class = scenario.get_technique_class() + scenario.set_params_from_args( + args={ + "objective_target": mock_objective_target, + "scenario_techniques": [technique_class("light")], + "include_baseline": False, + "max_attempts_per_objective": max_attempts, + } + ) + scenario._resolve_dataset_groups_for_estimate_async = AsyncMock(return_value=(groups, summaries)) + + with patch.object(scenario, "_get_attack_technique_factories", return_value=factories): + estimate = await scenario.get_run_size_estimate_async() + techniques = scenario._build_techniques_dict(objective_target=mock_objective_target) + + assert len(techniques) == 9 + assert {bundle.name for bundle in techniques.values()} == _LIGHT_TECHNIQUES + parsed_identifiers = [AdaptiveTechniqueIdentifier.parse(identifier) for identifier in techniques] + assert all(identifier is not None for identifier in parsed_identifiers) + assert len({identifier.factory_hash for identifier in parsed_identifiers if identifier is not None}) == 9 + assert len({identifier.technique_eval_hash for identifier in parsed_identifiers if identifier is not None}) == 9 + assert estimate.adaptive_details is not None + assert estimate.adaptive_details.selected_candidate_technique_count == 9 + assert estimate.adaptive_details.candidate_technique_count == 9 + assert estimate.adaptive_details.techniques_per_objective_upper_bound == max_attempts + assert estimate.adaptive_details.technique_attempt_count_upper_bound == expected_attempts + + @pytest.mark.parametrize(("aggregate_name", "expected_candidate_count"), [("light", 9), ("core", 14)]) + async def test_aggregate_attempt_bound_increases_until_distinct_candidate_count( + self, + mock_objective_target, + mock_objective_scorer, + aggregate_name, + expected_candidate_count, + ): + technique_class = TextAdaptive.get_technique_class() + aggregate = technique_class(aggregate_name) + selected_names = {technique.value for technique in technique_class.expand({aggregate})} + assert len(selected_names) == expected_candidate_count + + shared_attack_identifier = _mock_id("SharedAttackImplementation") + factories = { + name: _make_fake_factory( + attack_identifier=shared_attack_identifier, + factory_identifier=_mock_id(f"Factory_{name}"), + ) + for name in selected_names + } + groups = {"adaptive": [_make_seed_group(value=f"obj-{index}") for index in range(21)]} + summaries = [ + ScenarioDatasetSummary( + name="adaptive", + logical_seed_group_count=21, + selected_seed_group_count=21, + ) + ] + scenario = TextAdaptive(objective_scorer=mock_objective_scorer) + scenario._resolve_dataset_groups_for_estimate_async = AsyncMock(return_value=(groups, summaries)) + effective_caps: list[int] = [] + + with patch.object(scenario, "_get_attack_technique_factories", return_value=factories): + for limit in range(1, expected_candidate_count + 3): + scenario.set_params_from_args( + args={ + "objective_target": mock_objective_target, + "scenario_techniques": [aggregate], + "include_baseline": False, + "max_attempts_per_objective": limit, + } + ) + estimate = await scenario.get_run_size_estimate_async() + + assert estimate.adaptive_details is not None + expected_effective_cap = min(limit, expected_candidate_count) + effective_caps.append(estimate.adaptive_details.techniques_per_objective_upper_bound) + assert estimate.adaptive_details.candidate_technique_count == expected_candidate_count + assert estimate.adaptive_details.techniques_per_objective_upper_bound == expected_effective_cap + assert estimate.adaptive_details.technique_attempt_count_upper_bound == 21 * expected_effective_cap + + assert effective_caps == [ + *range(1, expected_candidate_count + 1), + expected_candidate_count, + expected_candidate_count, + ] + + def test_exact_duplicate_registered_technique_dedupes_only_itself( + self, + mock_objective_target, + mock_objective_scorer, + ): + scenario = TextAdaptive(objective_scorer=mock_objective_scorer) + technique_class = scenario.get_technique_class() + scenario._scenario_techniques = [ + technique_class("role_play_movie_script"), + technique_class("role_play_movie_script"), + ] + factory = _make_fake_factory() + + with patch.object( + scenario, + "_get_attack_technique_factories", + return_value={"role_play_movie_script": factory}, + ): + techniques = scenario._build_techniques_dict(objective_target=mock_objective_target) + + assert len(techniques) == 1 + factory.create.assert_called_once() + async def test_incompatible_seed_technique_is_filtered_per_objective( self, mock_objective_target, mock_objective_scorer ): @@ -692,11 +861,19 @@ async def test_baseline_emitted_at_index_zero_by_default(self, mock_objective_ta scenario must prepend a baseline atomic attack at index 0. """ groups = {"violence": [_make_seed_group(value="obj", harm_categories=["violence"])]} - with patch.object( - CompoundDatasetAttackConfiguration, - "get_attack_groups_by_dataset_async", - new_callable=AsyncMock, - return_value=groups, + with ( + patch.object( + CompoundDatasetAttackConfiguration, + "get_attack_groups_by_dataset_async", + new_callable=AsyncMock, + return_value=groups, + ), + patch.object( + CompoundDatasetAttackConfiguration, + "resolve_attack_groups_for_estimate_async", + new_callable=AsyncMock, + return_value=(groups, groups), + ), ): scenario = TextAdaptive(objective_scorer=mock_objective_scorer) with warnings.catch_warnings(): @@ -708,3 +885,9 @@ async def test_baseline_emitted_at_index_zero_by_default(self, mock_objective_ta assert scenario._atomic_attacks[0].atomic_attack_name == "baseline", ( f"baseline must be prepended at index 0; got {[a.atomic_attack_name for a in scenario._atomic_attacks]}" ) + estimate = await scenario.get_run_size_estimate_async() + plan = scenario._build_run_plan() + planned_units = sum(len(group.seed_group_ids) for group in plan.atomic_groups) + assert planned_units == 2 + assert estimate.total_attack_count == planned_units + assert [component.count for component in estimate.components] == [1, 1] diff --git a/tests/unit/scenario/test_default_run_size_estimates.py b/tests/unit/scenario/test_default_run_size_estimates.py index 88152579e9..f636185315 100644 --- a/tests/unit/scenario/test_default_run_size_estimates.py +++ b/tests/unit/scenario/test_default_run_size_estimates.py @@ -8,26 +8,26 @@ import pytest -from pyrit.executor.attack.core.attack_config import AttackScoringConfig +from pyrit.executor.attack import AttackScoringConfig from pyrit.models import ( AttackSeedGroup, AttackTechniqueSeedGroup, ComponentIdentifier, ScenarioDatasetSizeCap, ScenarioDatasetSummary, + ScenarioRunSizeEstimateCondition, SeedObjective, SeedPrompt, SeedSimulatedConversation, ) +from pyrit.models.catalog import ScenarioRunSizeEstimateStatus from pyrit.prompt_target import PromptTarget -from pyrit.scenario.core import BaselineAttackPolicy, DatasetAttackConfiguration, Scenario, ScenarioTechnique -from pyrit.scenario.scenarios.adaptive.text_adaptive import TextAdaptive -from pyrit.scenario.scenarios.airt.jailbreak import Jailbreak -from pyrit.scenario.scenarios.airt.psychosocial import Psychosocial -from pyrit.scenario.scenarios.benchmark.adversarial import AdversarialBenchmark -from pyrit.scenario.scenarios.foundry.red_team_agent import FoundryComposite, FoundryTechnique, RedTeamAgent -from pyrit.scenario.scenarios.garak.encoding import Encoding -from pyrit.scenario.scenarios.garak.web_injection import WebInjection +from pyrit.scenario import BaselineAttackPolicy, DatasetAttackConfiguration, Scenario, ScenarioTechnique +from pyrit.scenario.scenarios.adaptive import TextAdaptive +from pyrit.scenario.scenarios.airt import Jailbreak, Psychosocial +from pyrit.scenario.scenarios.benchmark import AdversarialBenchmark +from pyrit.scenario.scenarios.foundry import FoundryComposite, FoundryTechnique, RedTeamAgent +from pyrit.scenario.scenarios.garak import Encoding, WebInjection from pyrit.score import TrueFalseScorer @@ -455,8 +455,19 @@ async def test_adaptive_estimate_is_target_conditional_and_does_not_multiply_tec scenario._resolve_dataset_groups_for_estimate_async = AsyncMock(return_value=_resolved_groups({"adaptive": 3})) estimate = await scenario.get_default_run_size_estimate_async() - assert estimate.estimated_attack_count is None + + assert estimate.status is ScenarioRunSizeEstimateStatus.Conditional + assert estimate.total_attack_count is None + assert estimate.minimum_attack_count == 3 + assert estimate.maximum_attack_count == 6 assert [component.count for component in estimate.components] == [3, 3] + assert estimate.adaptive_details is not None + assert estimate.adaptive_details.objective_count == 3 + assert estimate.adaptive_details.selected_candidate_technique_count == 2 + assert estimate.adaptive_details.candidate_technique_count == 2 + assert estimate.adaptive_details.max_attempts_per_objective == 3 + assert estimate.adaptive_details.techniques_per_objective_upper_bound == 2 + assert estimate.adaptive_details.technique_attempt_count_upper_bound == 6 @pytest.mark.usefixtures("patch_central_database") @@ -486,13 +497,115 @@ async def test_adaptive_estimate_counts_exact_compatible_outer_envelopes_with_ta estimate = await scenario.get_run_size_estimate_async() assert estimate.estimated_attack_count == 2 assert [component.count for component in estimate.components] == [2] - assert "7 selected technique attempts" in estimate.note + assert "Up to 1 selected technique attempts" in estimate.note + assert estimate.adaptive_details is not None + assert estimate.adaptive_details.objective_count == 2 + assert estimate.adaptive_details.selected_candidate_technique_count == 2 + assert estimate.adaptive_details.candidate_technique_count == 1 + assert estimate.adaptive_details.techniques_per_objective_upper_bound == 1 + assert estimate.adaptive_details.technique_attempt_count_upper_bound == 2 scenario.set_params_from_args(args={"include_baseline": False}) estimate_without_target = await scenario.get_run_size_estimate_async() assert estimate_without_target.estimated_attack_count is None +@pytest.mark.usefixtures("patch_central_database") +async def test_adaptive_estimate_caps_attempts_below_candidate_pool() -> None: + """A lower configured max-attempt cap bounds each objective before pool size.""" + with patch.object(TextAdaptive, "get_technique_class", return_value=_TwoTechniqueDefault): + scenario = TextAdaptive(objective_scorer=_scorer()) + target = MagicMock(spec=PromptTarget) + scenario.set_params_from_args( + args={ + "objective_target": target, + "include_baseline": False, + "max_attempts_per_objective": 1, + } + ) + scenario._resolve_dataset_groups_for_estimate_async = AsyncMock(return_value=_resolved_groups({"adaptive": 3})) + dispatcher = MagicMock() + dispatcher.compatible_techniques.side_effect = [["one", "two"], ["one"], ["two"]] + + with ( + patch.object( + scenario, + "_build_techniques_dict", + return_value={"one": MagicMock(), "two": MagicMock()}, + ), + patch( + "pyrit.scenario.scenarios.adaptive.adaptive_scenario.AdaptiveTechniqueDispatcher", + return_value=dispatcher, + ), + ): + estimate = await scenario.get_run_size_estimate_async() + + assert estimate.adaptive_details is not None + assert estimate.adaptive_details.objective_count == 3 + assert estimate.adaptive_details.selected_candidate_technique_count == 2 + assert estimate.adaptive_details.candidate_technique_count == 2 + assert estimate.adaptive_details.max_attempts_per_objective == 1 + assert estimate.adaptive_details.techniques_per_objective_upper_bound == 1 + assert estimate.adaptive_details.technique_attempt_count_upper_bound == 3 + + +@pytest.mark.usefixtures("patch_central_database") +async def test_adaptive_conditional_attempt_bound_uses_launch_wide_objective_maximum() -> None: + """A sampled compatibility preview cannot understate a capped launch's attempt bound.""" + with patch.object(TextAdaptive, "get_technique_class", return_value=_TwoTechniqueDefault): + scenario = TextAdaptive(objective_scorer=_scorer()) + target = MagicMock(spec=PromptTarget) + scenario.set_params_from_args( + args={ + "objective_target": target, + "include_baseline": False, + "max_attempts_per_objective": 3, + } + ) + + async def resolve_groups() -> tuple[dict[str, list[AttackSeedGroup]], list[ScenarioDatasetSummary]]: + scenario._estimate_has_binding_size_cap = True + return _resolved_groups({"adaptive": 3}) + + scenario._resolve_dataset_groups_for_estimate_async = AsyncMock(side_effect=resolve_groups) + dispatcher = MagicMock() + dispatcher.compatible_techniques.side_effect = [["one"], [], []] + + with ( + patch.object( + scenario, + "_build_techniques_dict", + return_value={"one": MagicMock(), "two": MagicMock()}, + ), + patch( + "pyrit.scenario.scenarios.adaptive.adaptive_scenario.AdaptiveTechniqueDispatcher", + return_value=dispatcher, + ), + ): + estimate = await scenario.get_run_size_estimate_async() + + assert estimate.status is ScenarioRunSizeEstimateStatus.Conditional + assert estimate.minimum_attack_count is None + assert estimate.maximum_attack_count == 3 + assert estimate.components[0].count == 1 + assert estimate.adaptive_details is not None + assert estimate.adaptive_details.objective_count == 3 + assert estimate.adaptive_details.techniques_per_objective_upper_bound == 2 + assert estimate.adaptive_details.technique_attempt_count_upper_bound == 6 + + +@pytest.mark.usefixtures("patch_central_database") +async def test_adaptive_estimate_rejects_non_positive_attempt_limit_without_target() -> None: + """Invalid attempt limits fail explicitly before constructing estimate metadata.""" + with patch.object(TextAdaptive, "get_technique_class", return_value=_TwoTechniqueDefault): + scenario = TextAdaptive(objective_scorer=_scorer()) + scenario.set_params_from_args(args={"include_baseline": False, "max_attempts_per_objective": 0}) + scenario._resolve_dataset_groups_for_estimate_async = AsyncMock(return_value=_resolved_groups({"adaptive": 3})) + + with pytest.raises(ValueError, match="max_attempts_per_objective must be >= 1, got 0"): + await scenario.get_run_size_estimate_async() + + @pytest.mark.usefixtures("patch_central_database") async def test_jailbreak_estimate_exposes_template_attempt_and_target_capability_axes() -> None: """Jailbreak reports guaranteed inline work separately from conditional system delivery.""" @@ -501,7 +614,14 @@ async def test_jailbreak_estimate_exposes_template_attempt_and_target_capability scenario._resolve_dataset_groups_for_estimate_async = AsyncMock(return_value=_resolved_groups({"harmbench": 4})) estimate = await scenario.get_default_run_size_estimate_async() - assert estimate.estimated_attack_count is None + + assert estimate.status is ScenarioRunSizeEstimateStatus.Conditional + assert estimate.total_attack_count is None + assert estimate.minimum_attack_count == 12 + assert estimate.maximum_attack_count == 20 + assert estimate.components[2].condition is ScenarioRunSizeEstimateCondition.TargetCapabilities + assert estimate.model_dump(mode="json")["minimum_attack_count"] == 12 + assert estimate.model_dump(mode="json")["maximum_attack_count"] == 20 assert [component.count for component in estimate.components] == [4, 8, 8] assert "2 template(s) x 4 selected logical seed group(s) x 1 selected" in estimate.note assert "Baseline adds one unit per selected seed group (4 units)" in estimate.note From e5be6c7764b86e5ecda75eebcc8ae788f92673d0 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Thu, 20 Aug 2026 23:27:03 -0700 Subject: [PATCH 07/12] FIX: Cache conditional scenario estimates Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- pyrit/backend/services/scenario_service.py | 7 +++++- tests/unit/backend/test_scenario_service.py | 25 +++++++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/pyrit/backend/services/scenario_service.py b/pyrit/backend/services/scenario_service.py index eb0b7f9e63..cb3c4b70c3 100644 --- a/pyrit/backend/services/scenario_service.py +++ b/pyrit/backend/services/scenario_service.py @@ -19,6 +19,7 @@ ScenarioDefaultRunSizeEstimate, ScenarioRunSizeEstimate, ScenarioRunSizeEstimateRequest, + ScenarioRunSizeEstimateStatus, ) from pyrit.registry import ScenarioMetadata, ScenarioRegistry @@ -262,7 +263,11 @@ async def _compute_default_run_size_estimate_async( note=f"The scenario could not resolve its default inputs for estimation ({type(exc).__name__})." ) - expires_at = monotonic() + _UNAVAILABLE_CACHE_TTL_SECONDS if estimate.estimated_attack_count is None else None + expires_at = ( + monotonic() + _UNAVAILABLE_CACHE_TTL_SECONDS + if estimate.status is ScenarioRunSizeEstimateStatus.Unavailable + else None + ) cache = self._estimate_cache cache[cache_key] = (estimate, expires_at) cache.move_to_end(cache_key) diff --git a/tests/unit/backend/test_scenario_service.py b/tests/unit/backend/test_scenario_service.py index ceed2f32a8..84d6f12ddc 100644 --- a/tests/unit/backend/test_scenario_service.py +++ b/tests/unit/backend/test_scenario_service.py @@ -592,6 +592,31 @@ async def test_unavailable_estimate_cache_expires(self) -> None: assert second.default_run_size == estimate assert service._registry.create_instance.call_count == 2 + async def test_conditional_estimate_cache_does_not_expire(self) -> None: + """A successful conditional estimate remains cached despite having no exact total.""" + metadata = _make_scenario_metadata() + estimate = ScenarioDefaultRunSizeEstimate(status=ScenarioRunSizeEstimateStatus.Conditional) + scenario = MagicMock() + scenario.get_default_run_size_estimate_async = AsyncMock(return_value=estimate) + + with ( + patch.object(ScenarioService, "__init__", _initialize_test_service), + patch("pyrit.backend.services.scenario_service._UNAVAILABLE_CACHE_TTL_SECONDS", 0), + ): + service = ScenarioService() + service._registry = MagicMock() + service._registry.get_registered_class_metadata.return_value = metadata + service._registry.create_instance.return_value = scenario + + first = await service.get_scenario_async(scenario_name="test.scenario") + second = await service.get_scenario_async(scenario_name="test.scenario") + + assert first is not None + assert second is not None + assert first.default_run_size == estimate + assert second.default_run_size == estimate + assert service._registry.create_instance.call_count == 1 + async def test_estimate_cache_is_version_aware_and_bounded(self) -> None: """Scenario version changes invalidate estimates and the LRU stays bounded.""" estimate = ScenarioRunSizeEstimate( From f6518e97d9a7a11bb165e8b755c059295aa999e0 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Thu, 20 Aug 2026 23:45:00 -0700 Subject: [PATCH 08/12] FIX: Preserve rich scenario sizing after restack Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- pyrit/scenario/core/scenario.py | 33 ++++++++++++--- pyrit/scenario/scenarios/airt/psychosocial.py | 18 ++++++-- .../scenarios/benchmark/adversarial.py | 41 +++++++++++++++---- .../scenarios/foundry/red_team_agent.py | 18 ++++++-- pyrit/scenario/scenarios/garak/encoding.py | 22 ++++++++-- .../scenario/scenarios/garak/web_injection.py | 20 +++++++-- tests/unit/scenario/airt/test_jailbreak.py | 6 +-- 7 files changed, 124 insertions(+), 34 deletions(-) diff --git a/pyrit/scenario/core/scenario.py b/pyrit/scenario/core/scenario.py index 70be3461fd..877142a019 100644 --- a/pyrit/scenario/core/scenario.py +++ b/pyrit/scenario/core/scenario.py @@ -46,10 +46,14 @@ ScenarioRunPlanAtomicGroup, ScenarioRunPlanSeedGroup, ScenarioRunSizeComponent, - ScenarioRunSizeEstimate, ScenarioRunState, config_hash, ) +from pyrit.models.catalog import ( + ScenarioDefaultRunSizeEstimate, + ScenarioRunSizeEstimateStatus, + ScenarioRunSizeFactor, +) from pyrit.models.parameter import ComponentType, Parameter, RegistryReference from pyrit.prompt_target import PromptTarget from pyrit.prompt_target.common.target_requirements import TargetRequirements @@ -563,7 +567,7 @@ def _resolve_scenario_techniques(self, *, scenario_techniques: Any) -> list[Scen return self._technique_class.resolve(scenario_techniques, default=self._default_technique) @final - async def get_default_run_size_estimate_async(self) -> ScenarioRunSizeEstimate: + async def get_default_run_size_estimate_async(self) -> ScenarioDefaultRunSizeEstimate: """ Estimate the scenario's default planned execution units without starting a run. @@ -577,7 +581,9 @@ async def get_default_run_size_estimate_async(self) -> ScenarioRunSizeEstimate: return await self.get_run_size_estimate_async(target_is_configured=False) @final - async def get_run_size_estimate_async(self, *, target_is_configured: bool = False) -> ScenarioRunSizeEstimate: + async def get_run_size_estimate_async( + self, *, target_is_configured: bool = False + ) -> ScenarioDefaultRunSizeEstimate: """ Estimate the currently configured run without creating or persisting it. @@ -597,7 +603,7 @@ async def get_run_size_estimate_async(self, *, target_is_configured: bool = Fals self._estimate_target_is_configured = self._objective_target is not None return await self._estimate_run_size_async() - async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: + async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: """ Estimate a standard technique-by-seed-group scenario. @@ -618,6 +624,7 @@ async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: ScenarioRunSizeComponent( label="Baseline", count=seed_group_count, + factors=[ScenarioRunSizeFactor(label="selected logical seed groups", count=seed_group_count)], is_baseline=True, note="One unmodified prompt-sending unit per selected seed group.", ) @@ -643,8 +650,14 @@ async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: else: estimated_attack_count = None note += " The range covers every compatibility mix that the randomized per-dataset caps can select." - return ScenarioRunSizeEstimate( - estimated_attack_count=estimated_attack_count, + status = ( + ScenarioRunSizeEstimateStatus.Exact + if estimated_attack_count is not None + else ScenarioRunSizeEstimateStatus.Conditional + ) + return ScenarioDefaultRunSizeEstimate( + status=status, + total_attack_count=estimated_attack_count, minimum_attack_count=minimum_attack_count, maximum_attack_count=maximum_attack_count, components=components, @@ -670,6 +683,10 @@ def _build_technique_size_components( ScenarioRunSizeComponent( label="Default technique sweep", count=seed_group_count * technique_count, + factors=[ + ScenarioRunSizeFactor(label="selected logical seed groups", count=seed_group_count), + ScenarioRunSizeFactor(label="default concrete techniques", count=technique_count), + ], ) ] @@ -695,6 +712,10 @@ def _build_technique_size_components( ScenarioRunSizeComponent( label=technique.value, count=compatible_count, + factors=[ + ScenarioRunSizeFactor(label="selected concrete techniques", count=1), + ScenarioRunSizeFactor(label="compatible logical seed groups", count=compatible_count), + ], ) ) return components diff --git a/pyrit/scenario/scenarios/airt/psychosocial.py b/pyrit/scenario/scenarios/airt/psychosocial.py index 338931f643..834d02d26e 100644 --- a/pyrit/scenario/scenarios/airt/psychosocial.py +++ b/pyrit/scenario/scenarios/airt/psychosocial.py @@ -32,9 +32,13 @@ ) from pyrit.models import ( ScenarioRunSizeComponent, - ScenarioRunSizeEstimate, SeedPrompt, ) +from pyrit.models.catalog import ( + ScenarioDefaultRunSizeEstimate, + ScenarioRunSizeEstimateStatus, + ScenarioRunSizeFactor, +) from pyrit.models.parameter import Parameter from pyrit.prompt_normalizer.converter_configuration import ConverterConfiguration from pyrit.scenario.core.atomic_attack import AtomicAttack @@ -489,7 +493,7 @@ async def _resolve_seed_groups_by_dataset_async( self._dataset_config = rebuilt return await super()._resolve_seed_groups_by_dataset_async(apply_sampling=apply_sampling) - async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: + async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: """ Estimate the independent sub-harm technique sweeps and per-harm baselines. @@ -505,6 +509,10 @@ async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: ScenarioRunSizeComponent( label=f"{dataset_name} technique sweep", count=seed_group_count * technique_count, + factors=[ + ScenarioRunSizeFactor(label="selected logical seed groups", count=seed_group_count), + ScenarioRunSizeFactor(label="default concrete techniques", count=technique_count), + ], ) ) if self._include_baseline: @@ -512,12 +520,14 @@ async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: ScenarioRunSizeComponent( label=f"{dataset_name} baseline", count=seed_group_count, + factors=[ScenarioRunSizeFactor(label="selected logical seed groups", count=seed_group_count)], is_baseline=True, note="Psychosocial uses a distinct baseline and scorer for each sub-harm.", ) ) - return ScenarioRunSizeEstimate( - estimated_attack_count=sum(component.count for component in components), + return ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Exact, + total_attack_count=sum(component.count for component in components), components=components, datasets=datasets, note="Each default sub-harm is planned independently; retries and internal turns are excluded.", diff --git a/pyrit/scenario/scenarios/benchmark/adversarial.py b/pyrit/scenario/scenarios/benchmark/adversarial.py index 6d824ccee8..b2d9b4d622 100644 --- a/pyrit/scenario/scenarios/benchmark/adversarial.py +++ b/pyrit/scenario/scenarios/benchmark/adversarial.py @@ -17,7 +17,11 @@ ObjectiveTargetEvaluationIdentifier, ScenarioResult, ScenarioRunSizeComponent, - ScenarioRunSizeEstimate, +) +from pyrit.models.catalog import ( + ScenarioDefaultRunSizeEstimate, + ScenarioRunSizeEstimateStatus, + ScenarioRunSizeFactor, ) from pyrit.models.parameter import Parameter from pyrit.registry import AttackTechniqueRegistry, TargetRegistry @@ -203,7 +207,7 @@ def __init__( scenario_result_id=scenario_result_id, ) - async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: + async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: """ Estimate the target-by-technique matrix using execution compatibility. @@ -227,6 +231,10 @@ async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: ScenarioRunSizeComponent( label=technique.value, count=compatible_count, + factors=[ + ScenarioRunSizeFactor(label="selected concrete techniques", count=1), + ScenarioRunSizeFactor(label="compatible logical seed groups", count=compatible_count), + ], note="Count per adversarial target.", ) ) @@ -246,7 +254,8 @@ async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: per_target_maximum = sampled_per_target_count target_names = self.params.get("adversarial_targets") or [] if not target_names: - return ScenarioRunSizeEstimate( + return ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Conditional, minimum_attack_count=per_target_minimum, components=per_target_components, datasets=datasets, @@ -259,11 +268,22 @@ async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: resolved_targets = self._resolve_adversarial_targets(target_names=target_names) target_count = len(resolved_targets) components = [ - component.model_copy(update={"count": component.count * target_count, "note": None}) + component.model_copy( + update={ + "count": component.count * target_count, + "factors": [ + *component.factors[:1], + ScenarioRunSizeFactor(label="adversarial targets", count=target_count), + *component.factors[1:], + ], + "note": None, + } + ) for component in per_target_components ] if self._use_cached: - return ScenarioRunSizeEstimate( + return ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Conditional, minimum_attack_count=0, maximum_attack_count=per_target_maximum * target_count if per_target_maximum is not None else None, components=components, @@ -274,7 +294,8 @@ async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: ), ) if self._estimate_has_binding_size_cap and compatibility_bounds is None: - return ScenarioRunSizeEstimate( + return ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Conditional, components=components, datasets=datasets, note=( @@ -288,7 +309,8 @@ async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: and per_target_maximum is not None and per_target_minimum != per_target_maximum ): - return ScenarioRunSizeEstimate( + return ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Conditional, minimum_attack_count=per_target_minimum * target_count, maximum_attack_count=per_target_maximum * target_count, components=components, @@ -298,8 +320,9 @@ async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: "Baseline is forbidden." ), ) - return ScenarioRunSizeEstimate( - estimated_attack_count=sum(component.count for component in components), + return ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Exact, + total_attack_count=sum(component.count for component in components), components=components, datasets=datasets, note="Baseline is forbidden; retries and internal attack turns are excluded.", diff --git a/pyrit/scenario/scenarios/foundry/red_team_agent.py b/pyrit/scenario/scenarios/foundry/red_team_agent.py index 958683adc1..d65c2cec5c 100644 --- a/pyrit/scenario/scenarios/foundry/red_team_agent.py +++ b/pyrit/scenario/scenarios/foundry/red_team_agent.py @@ -53,7 +53,11 @@ from pyrit.models import ( AttackSeedGroup, ScenarioRunSizeComponent, - ScenarioRunSizeEstimate, +) +from pyrit.models.catalog import ( + ScenarioDefaultRunSizeEstimate, + ScenarioRunSizeEstimateStatus, + ScenarioRunSizeFactor, ) from pyrit.prompt_normalizer.converter_configuration import ConverterConfiguration from pyrit.prompt_target import PromptTarget @@ -418,7 +422,7 @@ def _resolve_foundry_techniques( self._scenario_composites = composites return flat - async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: + async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: """ Estimate one selected seed population per resolved Foundry composition. @@ -431,6 +435,10 @@ async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: ScenarioRunSizeComponent( label=composition.name, count=selected_count, + factors=[ + ScenarioRunSizeFactor(label="resolved Foundry composites", count=1), + ScenarioRunSizeFactor(label="selected logical seed groups", count=selected_count), + ], ) for composition in self._scenario_composites ] @@ -439,11 +447,13 @@ async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: ScenarioRunSizeComponent( label="Baseline", count=selected_count, + factors=[ScenarioRunSizeFactor(label="selected logical seed groups", count=selected_count)], is_baseline=True, ) ) - return ScenarioRunSizeEstimate( - estimated_attack_count=sum(component.count for component in components), + return ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Exact, + total_attack_count=sum(component.count for component in components), components=components, datasets=datasets, note="Counts one population per resolved Foundry composite, not per flattened constituent technique.", diff --git a/pyrit/scenario/scenarios/garak/encoding.py b/pyrit/scenario/scenarios/garak/encoding.py index 43bc565ea2..37b34f1573 100644 --- a/pyrit/scenario/scenarios/garak/encoding.py +++ b/pyrit/scenario/scenarios/garak/encoding.py @@ -27,11 +27,15 @@ from pyrit.models import ( AttackSeedGroup, ScenarioRunSizeComponent, - ScenarioRunSizeEstimate, Seed, SeedObjective, SeedPrompt, ) +from pyrit.models.catalog import ( + ScenarioDefaultRunSizeEstimate, + ScenarioRunSizeEstimateStatus, + ScenarioRunSizeFactor, +) from pyrit.prompt_normalizer.converter_configuration import ConverterConfiguration from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.attack_technique import AttackTechnique @@ -230,7 +234,7 @@ async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list atomic_attacks.extend(self._get_converter_attacks(context=context)) return atomic_attacks - async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: + async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: """ Estimate converter variants crossed with raw and decode-template prompt configurations. @@ -246,6 +250,14 @@ async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: ScenarioRunSizeComponent( label="Encoding converter variants", count=seed_group_count * variant_count * prompt_configuration_count, + factors=[ + ScenarioRunSizeFactor(label="selected logical seed groups", count=seed_group_count), + ScenarioRunSizeFactor(label="concrete converter variants", count=variant_count), + ScenarioRunSizeFactor( + label="raw plus decode prompt configurations", + count=prompt_configuration_count, + ), + ], note=( f"{seed_group_count} selected seed groups x {variant_count} concrete converter variants x " f"{prompt_configuration_count} prompt configurations. Base64 and ASCII85 each map to two " @@ -258,11 +270,13 @@ async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: ScenarioRunSizeComponent( label="Baseline", count=seed_group_count, + factors=[ScenarioRunSizeFactor(label="selected logical seed groups", count=seed_group_count)], is_baseline=True, ) ) - return ScenarioRunSizeEstimate( - estimated_attack_count=sum(component.count for component in components), + return ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Exact, + total_attack_count=sum(component.count for component in components), components=components, datasets=datasets, note=( diff --git a/pyrit/scenario/scenarios/garak/web_injection.py b/pyrit/scenario/scenarios/garak/web_injection.py index 5b6bf4e997..d03f85d154 100644 --- a/pyrit/scenario/scenarios/garak/web_injection.py +++ b/pyrit/scenario/scenarios/garak/web_injection.py @@ -16,10 +16,14 @@ AttackSeedGroup, ScenarioDatasetSummary, ScenarioRunSizeComponent, - ScenarioRunSizeEstimate, SeedObjective, SeedPrompt, ) +from pyrit.models.catalog import ( + ScenarioDefaultRunSizeEstimate, + ScenarioRunSizeEstimateStatus, + ScenarioRunSizeFactor, +) from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.attack_technique import AttackTechnique from pyrit.scenario.core.dataset_configuration import DatasetAttackConfiguration @@ -528,7 +532,7 @@ def _build_synthesized_seed_groups( ) return seed_groups_by_technique - async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: + async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: """ Estimate the technique-specific synthesized populations and their shared baseline. @@ -561,6 +565,7 @@ async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: ScenarioRunSizeComponent( label=f"{technique_name} synthesized prompts", count=len(seed_groups), + factors=[ScenarioRunSizeFactor(label="synthesized logical seed groups", count=len(seed_groups))], ) for technique_name, seed_groups in seed_groups_by_technique.items() ] @@ -570,12 +575,19 @@ async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: ScenarioRunSizeComponent( label="Baseline", count=synthesized_count, + factors=[ + ScenarioRunSizeFactor( + label="all synthesized logical seed groups", + count=synthesized_count, + ) + ], is_baseline=True, note="The baseline runs over the union of all default technique populations.", ) ) - return ScenarioRunSizeEstimate( - estimated_attack_count=sum(component.count for component in components), + return ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Exact, + total_attack_count=sum(component.count for component in components), components=components, datasets=datasets, note=( diff --git a/tests/unit/scenario/airt/test_jailbreak.py b/tests/unit/scenario/airt/test_jailbreak.py index 580a7c226b..e8db0b18c3 100644 --- a/tests/unit/scenario/airt/test_jailbreak.py +++ b/tests/unit/scenario/airt/test_jailbreak.py @@ -89,7 +89,7 @@ def mock_memory_seed_groups() -> list[AttackSeedGroup]: @pytest.fixture -def mock_objective_target() -> PromptTarget: +def mock_objective_target() -> MagicMock: """Create a mock objective target that cannot carry native system-prompt delivery. ``configuration.includes(...)`` returns ``False`` so the default technique set degrades to the @@ -102,7 +102,7 @@ def mock_objective_target() -> PromptTarget: @pytest.fixture -def mock_capable_target() -> PromptTarget: +def mock_capable_target() -> MagicMock: """Create a mock objective target that natively supports editable history + system prompts.""" mock = MagicMock(spec=PromptTarget) mock.get_identifier.return_value = ComponentIdentifier(class_name="MockCapableTarget", class_module="test") @@ -111,7 +111,7 @@ def mock_capable_target() -> PromptTarget: @pytest.fixture -def mock_objective_scorer() -> TrueFalseInverterScorer: +def mock_objective_scorer() -> MagicMock: """Create a mock scorer for testing.""" mock = MagicMock(spec=TrueFalseInverterScorer) mock.get_identifier.return_value = ComponentIdentifier(class_name="MockObjectiveScorer", class_module="test") From f8bac7729d0c3ea79c96fcaadf62033c34b5f614 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Thu, 20 Aug 2026 23:57:29 -0700 Subject: [PATCH 09/12] FIX: Clarify conditional scenario estimates Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../src/components/Scenarios/ScenarioRunEstimate.test.tsx | 5 +++-- frontend/src/components/Scenarios/ScenarioRunEstimate.tsx | 4 +++- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/frontend/src/components/Scenarios/ScenarioRunEstimate.test.tsx b/frontend/src/components/Scenarios/ScenarioRunEstimate.test.tsx index 4ecf198111..32d229e1df 100644 --- a/frontend/src/components/Scenarios/ScenarioRunEstimate.test.tsx +++ b/frontend/src/components/Scenarios/ScenarioRunEstimate.test.tsx @@ -706,8 +706,9 @@ describe('ScenarioRunEstimate', () => { /> , ) - expect(screen.getByText('Exact total')).toBeInTheDocument() - expect(screen.getByText('unavailable')).toBeInTheDocument() + expect(screen.getByText('Select targets')).toBeInTheDocument() + expect(screen.getByText('to calculate')).toBeInTheDocument() + expect(screen.queryByText('Exact total')).not.toBeInTheDocument() }) it('does not render implementation terminology in the shared estimate surfaces', () => { diff --git a/frontend/src/components/Scenarios/ScenarioRunEstimate.tsx b/frontend/src/components/Scenarios/ScenarioRunEstimate.tsx index f15d4d7583..eb21a7aee8 100644 --- a/frontend/src/components/Scenarios/ScenarioRunEstimate.tsx +++ b/frontend/src/components/Scenarios/ScenarioRunEstimate.tsx @@ -234,7 +234,9 @@ function resultOperand(estimate: ScenarioRunEstimate): CalculationOperand { result: true, } } - return { id: 'result', value: 'Exact total', label: 'unavailable', result: true } + return estimate.scope === 'default' + ? { id: 'result', value: 'Select targets', label: 'to calculate', result: true } + : { id: 'result', value: 'Confirmed', label: 'at launch', result: true } } function adaptivePlannedCalculation(estimate: ScenarioRunEstimate): RunCalculation { From b842fdabdf86ea473b3e5868360313f521629cf2 Mon Sep 17 00:00:00 2001 From: Roman Lutz Date: Mon, 24 Aug 2026 19:18:16 -0700 Subject: [PATCH 10/12] TEST: Keep auth bootstrap responsive in slow API E2E Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- frontend/e2e/api.spec.ts | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/frontend/e2e/api.spec.ts b/frontend/e2e/api.spec.ts index 9b5556ef29..2e87ad51fe 100644 --- a/frontend/e2e/api.spec.ts +++ b/frontend/e2e/api.spec.ts @@ -138,8 +138,21 @@ test.describe("Scenarios API", () => { test.describe("Error Handling", () => { test("should display UI when backend is slow", async ({ page }) => { - // Intercept and delay API calls + // Auth configuration must resolve before the application can render. await page.route("**/api/**", async (route) => { + if (new URL(route.request().url()).pathname === "/api/auth/config") { + await route.fulfill({ + status: 200, + contentType: "application/json", + body: JSON.stringify({ + clientId: "", + tenantId: "", + allowedGroupIds: "", + }), + }); + return; + } + await new Promise((resolve) => setTimeout(resolve, 2000)); await route.continue(); }); From 4f99222c45df8084408486f27459bbdc8b9c3d8e Mon Sep 17 00:00:00 2001 From: Roman Lutz Date: Mon, 24 Aug 2026 19:21:12 -0700 Subject: [PATCH 11/12] STYLE: Remove stale jailbreak test import Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../src/components/Scenarios/scenarioRunEstimateAdapter.ts | 4 ++-- tests/unit/scenario/airt/test_jailbreak.py | 1 - 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/frontend/src/components/Scenarios/scenarioRunEstimateAdapter.ts b/frontend/src/components/Scenarios/scenarioRunEstimateAdapter.ts index 63f3bed303..c8fce80d8e 100644 --- a/frontend/src/components/Scenarios/scenarioRunEstimateAdapter.ts +++ b/frontend/src/components/Scenarios/scenarioRunEstimateAdapter.ts @@ -105,7 +105,7 @@ export function mapScenarioRunEstimate( total, minimum: response.minimum_attack_count ?? null, maximum: response.maximum_attack_count ?? null, - condition: isRichEstimate ? response.condition : null, + condition: isRichEstimate ? response.condition ?? null : null, components: response.components.map((component) => { const id = nextStableId('component', component.label, componentOccurrences) return { @@ -114,7 +114,7 @@ export function mapScenarioRunEstimate( count: component.count, factors: mapFactors(id, component.factors ?? []), isBaseline: component.is_baseline, - condition: component.condition, + condition: component.condition ?? null, note: component.note, } }), diff --git a/tests/unit/scenario/airt/test_jailbreak.py b/tests/unit/scenario/airt/test_jailbreak.py index e8db0b18c3..488d987782 100644 --- a/tests/unit/scenario/airt/test_jailbreak.py +++ b/tests/unit/scenario/airt/test_jailbreak.py @@ -3,7 +3,6 @@ """Tests for the Jailbreak class.""" -import logging from typing import Any from unittest.mock import AsyncMock, MagicMock, patch From 88f163f2e2f8ef3d9d5903d5712db1848f5434db Mon Sep 17 00:00:00 2001 From: Behnam Ousat Date: Wed, 2 Sep 2026 10:50:06 -0700 Subject: [PATCH 12/12] FIX: Restore scenario dataset sizing contract Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 49641a15-4119-43bd-becc-9c43eef407f2 --- pyrit/scenario/core/scenario.py | 27 ++++++++++++++++--- tests/unit/models/test_scenario_catalog.py | 4 +-- tests/unit/scenario/core/test_scenario.py | 25 +++++++++++++++++ .../test_default_run_size_estimates.py | 1 + 4 files changed, 52 insertions(+), 5 deletions(-) diff --git a/pyrit/scenario/core/scenario.py b/pyrit/scenario/core/scenario.py index 877142a019..2d6888f187 100644 --- a/pyrit/scenario/core/scenario.py +++ b/pyrit/scenario/core/scenario.py @@ -15,7 +15,7 @@ from collections.abc import Mapping, Sequence from enum import Enum from pathlib import Path -from typing import TYPE_CHECKING, Any, ClassVar, final +from typing import TYPE_CHECKING, Any, ClassVar, Literal, final try: # Built-in on Python 3.11+. Fall back to the ``exceptiongroup`` backport on 3.10 @@ -60,7 +60,11 @@ from pyrit.registry import ScorerRegistry from pyrit.registry.resolution import resolve_declared_params, resolve_reference_value from pyrit.scenario.core.atomic_attack import AtomicAttack -from pyrit.scenario.core.dataset_configuration import DatasetAttackConfiguration, read_only_dataset_resolution +from pyrit.scenario.core.dataset_configuration import ( + CompoundDatasetAttackConfiguration, + DatasetAttackConfiguration, + read_only_dataset_resolution, +) from pyrit.scenario.core.scenario_context import ScenarioContext from pyrit.scenario.core.scenario_target_defaults import get_default_scorer_target from pyrit.scenario.core.scenario_technique import ScenarioTechnique @@ -140,6 +144,10 @@ class Scenario(ABC): #: Whether the default estimator must mirror matrix-builder seed compatibility. RUN_SIZE_USES_FACTORY_COMPATIBILITY: ClassVar[bool] = False + #: How a generic dataset-size run override is interpreted. ``None`` derives the + #: standard behavior from the default configuration. + DATASET_SIZE_LIMIT_OVERRIDE_SCOPE: ClassVar[Literal["per_dataset", "combined", "unsupported"] | None] = None + def __init_subclass__(cls, **kwargs: Any) -> None: """ Enforce the keyword-only constructor contract on subclasses. @@ -263,6 +271,19 @@ def __init__( # before _build_atomic_attacks_async is awaited so overrides can read it. self._include_baseline: bool = False + def get_dataset_size_limit_override_scope(self) -> Literal["per_dataset", "combined", "unsupported"]: + """ + Return how this scenario interprets a generic dataset-size run override. + + Returns: + Literal: The explicit override scope exposed through the scenario catalog. + """ + if self.DATASET_SIZE_LIMIT_OVERRIDE_SCOPE is not None: + return self.DATASET_SIZE_LIMIT_OVERRIDE_SCOPE + if isinstance(self._default_dataset_config, CompoundDatasetAttackConfiguration): + return "per_dataset" + return "per_dataset" if len(self._default_dataset_config.dataset_names) <= 1 else "combined" + @property def name(self) -> str: """The name of the scenario.""" @@ -831,7 +852,7 @@ async def _resolve_dataset_groups_for_estimate_async( selected_count = len(selected_groups.get(name, [])) selection_note = None if selected_count != logical_count: - selection_note = f"The default selection uses {selected_count} of {logical_count} logical seed groups." + selection_note = f"The default selection uses {selected_count} of {logical_count} available objectives." datasets.append( ScenarioDatasetSummary( name=name, diff --git a/tests/unit/models/test_scenario_catalog.py b/tests/unit/models/test_scenario_catalog.py index 9943beb7ff..60000fbee4 100644 --- a/tests/unit/models/test_scenario_catalog.py +++ b/tests/unit/models/test_scenario_catalog.py @@ -275,7 +275,7 @@ def test_conditional_estimate_exposes_dataset_counts_structurally() -> None: name="harmbench", logical_seed_group_count=100, selected_seed_group_count=4, - selection_note="The default selection uses 4 of 100 logical seed groups.", + selection_note="The default selection uses 4 of 100 available objectives.", configured_caps=[ ScenarioDatasetSizeCap( label="per-dataset cap", @@ -301,7 +301,7 @@ def test_conditional_estimate_exposes_dataset_counts_structurally() -> None: "kind": "dataset", "logical_seed_group_count": 100, "selected_seed_group_count": 4, - "selection_note": "The default selection uses 4 of 100 logical seed groups.", + "selection_note": "The default selection uses 4 of 100 available objectives.", "configured_caps": [ { "label": "per-dataset cap", diff --git a/tests/unit/scenario/core/test_scenario.py b/tests/unit/scenario/core/test_scenario.py index 23644d02ca..2cf7b7abd4 100644 --- a/tests/unit/scenario/core/test_scenario.py +++ b/tests/unit/scenario/core/test_scenario.py @@ -28,6 +28,7 @@ ) from pyrit.prompt_target import PromptTarget from pyrit.scenario import ( + CompoundDatasetAttackConfiguration, DatasetAttackConfiguration, DatasetConfiguration, ScenarioIdentifier, @@ -220,6 +221,30 @@ def test_subclass_implementing_build_atomic_attacks_async_is_concrete(): assert not ConcreteScenario.__abstractmethods__ +@pytest.mark.parametrize( + ("default_dataset_config", "expected_scope"), + [ + (DatasetAttackConfiguration(dataset_names=["one"]), "per_dataset"), + (DatasetAttackConfiguration(dataset_names=["one", "two"]), "combined"), + ( + CompoundDatasetAttackConfiguration.per_dataset( + dataset_names=["one", "two"], + max_dataset_size=4, + ), + "per_dataset", + ), + ], +) +def test_dataset_size_limit_override_scope_follows_default_configuration( + default_dataset_config: DatasetAttackConfiguration, + expected_scope: str, +) -> None: + """Dataset-size overrides preserve the scenario's default configuration semantics.""" + scenario = ConcreteScenario(version=1, default_dataset_config=default_dataset_config) + + assert scenario.get_dataset_size_limit_override_scope() == expected_scope + + @pytest.mark.usefixtures("patch_central_database") class TestScenarioInitialization: """Tests for Scenario class initialization.""" diff --git a/tests/unit/scenario/test_default_run_size_estimates.py b/tests/unit/scenario/test_default_run_size_estimates.py index f636185315..614211f1ca 100644 --- a/tests/unit/scenario/test_default_run_size_estimates.py +++ b/tests/unit/scenario/test_default_run_size_estimates.py @@ -135,6 +135,7 @@ async def test_ordinary_matrix_estimate_uses_planned_seed_units_and_baseline() - assert [component.count for component in estimate.components] == [4, 2] assert estimate.datasets[0].logical_seed_group_count == 3 assert estimate.datasets[0].selected_seed_group_count == 2 + assert estimate.datasets[0].selection_note == "The default selection uses 2 of 3 available objectives." @pytest.mark.usefixtures("patch_central_database")