From 948fcc88098ad4643ca581745bd52b664dbd98ca Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Fri, 23 Oct 2026 09:01:43 -0400 Subject: [PATCH 1/7] FEAT: Add ScenarioSuite registry support Merge code-registered read-only suites with stored editable definitions and surface suite origin across the API and UI. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 2403a4a9-6af6-46b2-b371-46de73a018fd --- .pyrit_conf_example | 8 +- doc/getting_started/pyrit_conf.md | 4 +- .../ScenarioSuiteEditor.test.tsx | 22 ++ .../ScenarioSuites/ScenarioSuiteEditor.tsx | 62 ++-- .../ScenarioSuiteLibrary.test.tsx | 25 ++ .../ScenarioSuites/ScenarioSuiteLibrary.tsx | 12 +- frontend/src/services/api.test.ts | 2 +- frontend/src/types/index.ts | 9 +- pyrit/backend/models/scenario_suites.py | 27 +- pyrit/backend/routes/scenario_suites.py | 38 ++- pyrit/backend/services/runtime_lifecycle.py | 11 +- .../services/scenario_suite_service.py | 72 ++--- pyrit/models/catalog/__init__.py | 8 +- pyrit/models/catalog/scenario_suite.py | 46 ++- pyrit/registry/__init__.py | 6 + pyrit/registry/components/__init__.py | 8 + .../components/scenario_suite_registry.py | 276 +++++++++++++++++ tests/unit/backend/test_runtime_lifecycle.py | 7 +- .../backend/test_scenario_suite_routes.py | 68 +++- .../backend/test_scenario_suite_service.py | 34 +- tests/unit/models/test_scenario_suite.py | 34 +- .../registry/test_scenario_suite_registry.py | 291 ++++++++++++++++++ 22 files changed, 962 insertions(+), 108 deletions(-) create mode 100644 pyrit/registry/components/scenario_suite_registry.py create mode 100644 tests/unit/registry/test_scenario_suite_registry.py diff --git a/.pyrit_conf_example b/.pyrit_conf_example index 7ad3514528..0a1515872f 100644 --- a/.pyrit_conf_example +++ b/.pyrit_conf_example @@ -149,9 +149,11 @@ allow_custom_initializers: false # otherwise DefaultAzureCredential is used. Defaults to ~/.pyrit/scenario_presets. # scenario_presets_source: https://account.blob.core.windows.net/container/scenario_presets -# Optional storage for ScenarioSuite definitions, which contain ordered references -# to scenario presets. This supports the same local directory and Azure Blob -# sources as scenario_presets_source. Defaults to ~/.pyrit/scenario_suites. +# Optional storage for operator-authored ScenarioSuite definitions, which contain +# ordered references to scenario presets. PyRIT merges these editable documents +# with read-only suites registered in code; a stored document takes precedence on +# a name collision. This supports the same local directory and Azure Blob sources +# as scenario_presets_source. Defaults to ~/.pyrit/scenario_suites. # scenario_suites_source: https://account.blob.core.windows.net/container/scenario_suites # Upload Folder for API-Created Targets diff --git a/doc/getting_started/pyrit_conf.md b/doc/getting_started/pyrit_conf.md index 0b4316d9a8..7d11c274bd 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -184,7 +184,9 @@ Stores ScenarioSuite definitions in a local directory or Azure Blob container. A scenario_suites_source: https://account.blob.core.windows.net/pyrit-storage/scenario-suites ``` -Suites store ordered references to scenario presets and resolve each preset's current status when read. A SAS query string may be included; otherwise, PyRIT uses `DefaultAzureCredential`. The default is `~/.pyrit/scenario_suites`. Local writes use the same hidden lock-file and atomic replacement behavior as scenario presets. +Suites store ordered references to scenario presets and resolve each preset's current status when read. The `ScenarioSuiteRegistry` merges editable suites from this source with read-only suites registered by code. A stored document takes precedence when both layers use the same name, while creating or updating a document under a registered name is rejected so one name has one owner. Registered suites report no document version. + +A SAS query string may be included; otherwise, PyRIT uses `DefaultAzureCredential`. The default is `~/.pyrit/scenario_suites`. Local writes use the same hidden lock-file and atomic replacement behavior as scenario presets. Startup scans the configured source so malformed suite documents are reported immediately without preventing the backend from starting. ### `initialization_scripts` diff --git a/frontend/src/components/ScenarioSuites/ScenarioSuiteEditor.test.tsx b/frontend/src/components/ScenarioSuites/ScenarioSuiteEditor.test.tsx index f8082bc296..a0ed421bca 100644 --- a/frontend/src/components/ScenarioSuites/ScenarioSuiteEditor.test.tsx +++ b/frontend/src/components/ScenarioSuites/ScenarioSuiteEditor.test.tsx @@ -68,6 +68,7 @@ function makeSuite(): ScenarioSuiteResponse { { entry_id: SECOND_ENTRY_ID, preset_name: 'retired_probe' }, ], }, + origin: 'stored', version: 'suite-v1', entry_statuses: [ { @@ -272,6 +273,27 @@ describe('ScenarioSuiteEditor', () => { )) }) + it('shows a registered suite as read-only and never offers mutation controls', async () => { + mockGetSuite.mockResolvedValue({ + ...makeSuite(), + origin: 'registered', + version: null, + }) + renderEditor('edit') + + expect(await screen.findByRole('heading', { name: 'View nightly_suite' })).toBeInTheDocument() + expect(screen.getByText( + 'This suite is defined in code. It is read-only and cannot be edited or deleted here.', + )).toBeInTheDocument() + expect(screen.getByTestId('suite-description-input')).toBeDisabled() + expect(screen.getByTestId('suite-preset-select')).toBeDisabled() + expect(screen.getByRole('button', { name: 'Move entry 2 up' })).toBeDisabled() + expect(screen.getByRole('button', { name: 'Remove entry 1' })).toBeDisabled() + expect(screen.queryByTestId('save-suite-btn')).not.toBeInTheDocument() + expect(screen.getByRole('button', { name: 'Back to suites' })).toBeEnabled() + expect(mockUpdateSuite).not.toHaveBeenCalled() + }) + it('removes entries and adds a registered preset without changing existing IDs', async () => { const user = userEvent.setup() renderEditor('edit') diff --git a/frontend/src/components/ScenarioSuites/ScenarioSuiteEditor.tsx b/frontend/src/components/ScenarioSuites/ScenarioSuiteEditor.tsx index 718756d0fd..e6a0e1b699 100644 --- a/frontend/src/components/ScenarioSuites/ScenarioSuiteEditor.tsx +++ b/frontend/src/components/ScenarioSuites/ScenarioSuiteEditor.tsx @@ -110,7 +110,9 @@ function ScenarioSuiteEditorContent({ mode, suiteName }: ScenarioSuiteEditorCont const [saveError, setSaveError] = useState(null) const currentDraft: SuiteDraft = { name, description, entries } + const readOnly = loaded?.origin === 'registered' const hasUnsavedChanges = status === 'ready' + && !readOnly && draftSignature(currentDraft) !== draftSignature(savedDraft) const shouldBlockNavigation = useCallback( (): boolean => hasUnsavedChanges && !allowNavigation.current, @@ -276,6 +278,10 @@ function ScenarioSuiteEditorContent({ mode, suiteName }: ScenarioSuiteEditorCont try { if (loaded === null) { await scenarioSuitesApi.create(suite) + } else if (loaded.version === null) { + setSaveError(`Suite '${loaded.suite.name}' is defined in code and cannot be edited here.`) + setSaving(false) + return } else { await scenarioSuitesApi.update(loaded.suite.name, suite, loaded.version) } @@ -336,7 +342,7 @@ function ScenarioSuiteEditorContent({ mode, suiteName }: ScenarioSuiteEditorCont >
- {editing ? `Edit ${decodedName}` : 'New suite'} + {readOnly ? `View ${decodedName}` : editing ? `Edit ${decodedName}` : 'New suite'} Arrange scenario presets in the order they should appear. Saving a suite does not run it. @@ -344,11 +350,20 @@ function ScenarioSuiteEditorContent({ mode, suiteName }: ScenarioSuiteEditorCont
void handleSubmit(event)} noValidate> + {readOnly && ( + + + This suite is defined in code. It is read-only and cannot be edited or deleted here. + + + )} {presetLoadError && ( - Preset library unavailable. Existing entries can still be edited and saved, but new presets cannot - be added. {presetLoadError} + Preset library unavailable. {readOnly + ? 'Existing entries remain visible, but current preset details may be unavailable.' + : 'Existing entries can still be edited and saved, but new presets cannot be added.'}{' '} + {presetLoadError} )} @@ -371,6 +386,7 @@ function ScenarioSuiteEditorContent({ mode, suiteName }: ScenarioSuiteEditorCont Schema version {loaded?.suite.schema_version ?? 1} {editing ? ` · Author: ${loaded?.suite.author ?? '—'}` : ''} + {readOnly ? ' · Defined in code' : ''} setDescription(data.value)} data-testid="suite-description-input" /> @@ -416,8 +432,10 @@ function ScenarioSuiteEditorContent({ mode, suiteName }: ScenarioSuiteEditorCont {availablePresets.length === 0 ? ( - No presets are available to add. Existing unavailable references can still be - preserved, reordered, removed, and saved. + {readOnly + ? 'No current preset definitions are available. The registered suite entries remain read-only.' + : 'No presets are available to add. Existing unavailable references can still be preserved, ' + + 'reordered, removed, and saved.'} ) : ( @@ -427,7 +445,7 @@ function ScenarioSuiteEditorContent({ mode, suiteName }: ScenarioSuiteEditorCont value={selectedPresetName} selectedOptions={selectedPresetName === '' ? [] : [selectedPresetName]} placeholder="Select a preset" - disabled={saving} + disabled={saving || readOnly} onOptionSelect={(_, data) => setSelectedPresetName(data.optionValue ?? '')} data-testid="suite-preset-select" > @@ -443,7 +461,7 @@ function ScenarioSuiteEditorContent({ mode, suiteName }: ScenarioSuiteEditorCont type="button" appearance="secondary" icon={} - disabled={saving || selectedPresetName === ''} + disabled={saving || readOnly || selectedPresetName === ''} onClick={addEntry} data-testid="add-suite-entry" > @@ -484,7 +502,7 @@ function ScenarioSuiteEditorContent({ mode, suiteName }: ScenarioSuiteEditorCont size="small" icon={} aria-label={`Move entry ${index + 1} up`} - disabled={saving || index === 0} + disabled={saving || readOnly || index === 0} onClick={() => moveEntry(entry.entry_id, -1)} /> + {!readOnly && ( + + )} diff --git a/frontend/src/components/ScenarioSuites/ScenarioSuiteLibrary.test.tsx b/frontend/src/components/ScenarioSuites/ScenarioSuiteLibrary.test.tsx index e68a5af085..a05fa243f6 100644 --- a/frontend/src/components/ScenarioSuites/ScenarioSuiteLibrary.test.tsx +++ b/frontend/src/components/ScenarioSuites/ScenarioSuiteLibrary.test.tsx @@ -39,6 +39,7 @@ function makeSuite(): ScenarioSuiteResponse { { entry_id: SECOND_ENTRY_ID, preset_name: 'nightly_probe' }, ], }, + origin: 'stored', version: 'suite-v1', entry_statuses: [ { @@ -59,6 +60,14 @@ function makeSuite(): ScenarioSuiteResponse { } } +function makeRegisteredSuite(): ScenarioSuiteResponse { + return { + ...makeSuite(), + origin: 'registered', + version: null, + } +} + function renderLibrary() { return render( @@ -132,6 +141,22 @@ describe('ScenarioSuiteLibrary', () => { expect(await screen.findByTestId('suite-empty-state')).toBeInTheDocument() }) + it('identifies registered suites and exposes only a read-only view action', async () => { + const user = userEvent.setup() + mockList.mockResolvedValue({ source: 'suites', items: [makeRegisteredSuite()] }) + + renderLibrary() + + const row = await screen.findByTestId('suite-row-nightly_suite') + expect(within(row).getByText('Defined in code')).toBeInTheDocument() + expect(within(row).getByRole('button', { name: 'View' })).toBeEnabled() + expect(within(row).getByRole('button', { name: 'Delete' })).toBeDisabled() + + await user.click(within(row).getByRole('button', { name: 'View' })) + expect(mockNavigate).toHaveBeenCalledWith('/registry/scenario-suites/nightly_suite/edit') + expect(mockRemove).not.toHaveBeenCalled() + }) + it('reports load failures and retries', async () => { const user = userEvent.setup() mockList.mockRejectedValueOnce(new Error('suite storage unavailable')) diff --git a/frontend/src/components/ScenarioSuites/ScenarioSuiteLibrary.tsx b/frontend/src/components/ScenarioSuites/ScenarioSuiteLibrary.tsx index 3ba6cf979d..edd0fafe00 100644 --- a/frontend/src/components/ScenarioSuites/ScenarioSuiteLibrary.tsx +++ b/frontend/src/components/ScenarioSuites/ScenarioSuiteLibrary.tsx @@ -30,6 +30,8 @@ import type { ScenarioSuiteResponse } from '@/types' import { useScenarioSuiteLibraryStyles } from './ScenarioSuiteLibrary.styles' import { NEW_SUITE_ROUTE, suiteEditorRoutePath } from './suiteRoutes' +const DEFINED_IN_CODE_HINT = 'Defined in code, so it cannot be edited or deleted here.' + function issueCount(item: ScenarioSuiteResponse): number { return item.entry_statuses.reduce( (count, status) => count + status.issues.length, @@ -188,6 +190,7 @@ export default function ScenarioSuiteLibrary() { {items.map((item) => { const issues = issueCount(item) + const editable = item.origin === 'stored' return ( @@ -200,6 +203,11 @@ export default function ScenarioSuiteLibrary() { Needs attention )} + {!editable && ( + + Defined in code + + )} {item.suite.description && ( @@ -247,12 +255,14 @@ export default function ScenarioSuiteLibrary() { onClick={() => navigate(suiteEditorRoutePath(item.suite.name))} data-testid={`edit-suite-${item.suite.name}`} > - Edit + {editable ? 'Edit' : 'View'} + + + + + + + ) +} diff --git a/frontend/src/components/ScenarioSuites/ScenarioSuiteExecutionPage.styles.ts b/frontend/src/components/ScenarioSuites/ScenarioSuiteExecutionPage.styles.ts new file mode 100644 index 0000000000..16c453ac36 --- /dev/null +++ b/frontend/src/components/ScenarioSuites/ScenarioSuiteExecutionPage.styles.ts @@ -0,0 +1,184 @@ +import { makeStyles, tokens } from '@fluentui/react-components' + +import { + MINIMUM_TOUCH_TARGET_SIZE, + NARROW_VIEWPORT_QUERY, + mobileTouchTarget, +} from '@/styles/touchTargets' +import { WORKSPACE_CANVAS_BACKGROUND } from '@/styles/workspaceBackground' + +export const useScenarioSuiteExecutionPageStyles = makeStyles({ + root: { + display: 'flex', + flexDirection: 'column', + width: '100%', + height: '100%', + minWidth: 0, + overflowY: 'auto', + overflowX: 'hidden', + backgroundColor: WORKSPACE_CANVAS_BACKGROUND, + }, + content: { + display: 'flex', + flexDirection: 'column', + width: '100%', + maxWidth: '96rem', + gap: tokens.spacingVerticalXL, + padding: tokens.spacingVerticalXXL, + marginInline: 'auto', + [NARROW_VIEWPORT_QUERY]: { + padding: `${tokens.spacingVerticalL} ${tokens.spacingHorizontalM}`, + gap: tokens.spacingVerticalL, + }, + }, + backLink: { + display: 'inline-flex', + alignItems: 'center', + alignSelf: 'flex-start', + gap: tokens.spacingHorizontalXS, + minHeight: MINIMUM_TOUCH_TARGET_SIZE, + color: tokens.colorBrandForegroundLink, + textDecorationLine: 'none', + ':hover': { + textDecorationLine: 'underline', + }, + ':focus-visible': { + outline: `2px solid ${tokens.colorStrokeFocus2}`, + outlineOffset: '2px', + }, + }, + header: { + display: 'flex', + alignItems: 'flex-start', + justifyContent: 'space-between', + gap: tokens.spacingHorizontalXL, + [NARROW_VIEWPORT_QUERY]: { + flexDirection: 'column', + alignItems: 'stretch', + }, + }, + headerIdentity: { + display: 'flex', + flexDirection: 'column', + minWidth: 0, + gap: tokens.spacingVerticalXS, + }, + titleRow: { + display: 'flex', + alignItems: 'center', + flexWrap: 'wrap', + gap: tokens.spacingHorizontalS, + }, + executionId: { + color: tokens.colorNeutralForeground3, + overflowWrap: 'anywhere', + }, + headerActions: { + display: 'flex', + flexShrink: 0, + gap: tokens.spacingHorizontalS, + [NARROW_VIEWPORT_QUERY]: { + width: '100%', + }, + }, + touchTarget: { + ...mobileTouchTarget, + }, + wideButton: { + [NARROW_VIEWPORT_QUERY]: { + flexGrow: 1, + }, + }, + metadata: { + display: 'grid', + gridTemplateColumns: 'repeat(3, minmax(10rem, 1fr))', + gap: `${tokens.spacingVerticalS} ${tokens.spacingHorizontalXL}`, + paddingTop: tokens.spacingVerticalM, + borderTop: `1px solid ${tokens.colorNeutralStroke2}`, + [NARROW_VIEWPORT_QUERY]: { + gridTemplateColumns: '1fr', + }, + }, + metadataItem: { + display: 'flex', + flexDirection: 'column', + gap: tokens.spacingVerticalXXS, + minWidth: 0, + }, + metadataLabel: { + color: tokens.colorNeutralForeground3, + }, + section: { + display: 'flex', + flexDirection: 'column', + gap: tokens.spacingVerticalM, + }, + sectionHeading: { + display: 'flex', + alignItems: 'baseline', + justifyContent: 'space-between', + flexWrap: 'wrap', + gap: tokens.spacingHorizontalM, + }, + sectionHint: { + color: tokens.colorNeutralForeground3, + }, + summaryGrid: { + display: 'grid', + gridTemplateColumns: 'repeat(auto-fit, minmax(10rem, 1fr))', + gap: tokens.spacingHorizontalM, + }, + summaryItem: { + display: 'flex', + flexDirection: 'column', + gap: tokens.spacingVerticalXXS, + padding: tokens.spacingVerticalL, + border: `1px solid ${tokens.colorNeutralStroke2}`, + borderRadius: tokens.borderRadiusLarge, + backgroundColor: tokens.colorNeutralBackground1, + }, + summaryValue: { + fontVariantNumeric: 'tabular-nums', + }, + tableScroll: { + width: '100%', + overflowX: 'auto', + border: `1px solid ${tokens.colorNeutralStroke2}`, + borderRadius: tokens.borderRadiusLarge, + backgroundColor: tokens.colorNeutralBackground1, + }, + table: { + minWidth: '62rem', + }, + entryIdentity: { + display: 'flex', + flexDirection: 'column', + gap: tokens.spacingVerticalXXS, + }, + entryMeta: { + color: tokens.colorNeutralForeground3, + }, + errorText: { + color: tokens.colorPaletteRedForeground1, + overflowWrap: 'anywhere', + }, + centeredState: { + display: 'flex', + flexDirection: 'column', + alignItems: 'center', + justifyContent: 'center', + gap: tokens.spacingVerticalM, + minHeight: '18rem', + textAlign: 'center', + }, + dialogContent: { + display: 'flex', + flexDirection: 'column', + gap: tokens.spacingVerticalM, + overflowWrap: 'anywhere', + }, + liveStatus: { + minHeight: '1.25rem', + color: tokens.colorNeutralForeground3, + }, +}) diff --git a/frontend/src/components/ScenarioSuites/ScenarioSuiteExecutionPage.test.tsx b/frontend/src/components/ScenarioSuites/ScenarioSuiteExecutionPage.test.tsx new file mode 100644 index 0000000000..fea7c27334 --- /dev/null +++ b/frontend/src/components/ScenarioSuites/ScenarioSuiteExecutionPage.test.tsx @@ -0,0 +1,246 @@ +import { render, screen, waitFor, within } from '@testing-library/react' +import userEvent from '@testing-library/user-event' +import { FluentProvider, webLightTheme } from '@fluentui/react-components' +import { MemoryRouter, Route, Routes } from 'react-router' + +import { + useScenarioSuiteExecution, + type UseScenarioSuiteExecutionResult, +} from '@/hooks/useScenarioSuiteExecution' +import { scenarioSuitesApi } from '@/services/api' +import type { + ScenarioRunSummary, + ScenarioSuiteExecutionEntrySummary, + ScenarioSuiteExecutionSummary, +} from '@/types' + +import ScenarioSuiteExecutionPage from './ScenarioSuiteExecutionPage' + +jest.mock('@/hooks/useScenarioSuiteExecution', () => ({ + useScenarioSuiteExecution: jest.fn(), + isTerminalSuiteExecutionStatus: (status: string) => ( + status === 'COMPLETED' || status === 'FAILED' || status === 'CANCELLED' + ), +})) + +jest.mock('@/services/api', () => ({ + scenarioSuitesApi: { + cancelExecution: jest.fn(), + }, +})) + +const mockUseScenarioSuiteExecution = useScenarioSuiteExecution as jest.Mock +const mockCancelExecution = scenarioSuitesApi.cancelExecution as jest.Mock +const mockRetry = jest.fn() +const mockApplyExecution = jest.fn() + +const FIRST_ENTRY_ID = '11111111-1111-4111-8111-111111111111' +const SECOND_ENTRY_ID = '22222222-2222-4222-8222-222222222222' + +function makeRun({ + runId, + status = 'IN_PROGRESS', + completedAttacks = 2, +}: { + runId: string + status?: ScenarioRunSummary['status'] + completedAttacks?: number +}): ScenarioRunSummary { + return { + scenario_result_id: runId, + scenario_name: 'scenario.example', + scenario_version: 1, + status, + created_at: '2026-10-23T12:00:00Z', + updated_at: '2026-10-23T12:01:00Z', + techniques_used: ['crescendo'], + total_attacks: 10, + completed_attacks: completedAttacks, + objective_achieved_rate: 0, + failed_attacks: [], + attack_retries: [], + total_retries: 0, + labels: {}, + } +} + +function makeEntry({ + entryId, + position, + presetName, + state = 'IN_PROGRESS', + run = null, + error = null, +}: { + entryId: string + position: number + presetName: string + state?: ScenarioSuiteExecutionEntrySummary['state'] + run?: ScenarioRunSummary | null + error?: string | null +}): ScenarioSuiteExecutionEntrySummary { + return { + entry_id: entryId, + entry_position: position, + preset_name: presetName, + preset_origin: 'stored', + preset_version: 'preset-v1', + state, + run, + error, + error_type: error ? 'ValueError' : null, + } +} + +function makeExecution( + overrides: Partial = {}, +): ScenarioSuiteExecutionSummary { + return { + suite_execution_id: 'execution-1', + suite_name: 'nightly_suite', + suite_origin: 'stored', + suite_version: 'suite-v1', + status: 'IN_PROGRESS', + created_at: '2026-10-23T12:00:00Z', + updated_at: '2026-10-23T12:01:00Z', + entries: [ + makeEntry({ + entryId: FIRST_ENTRY_ID, + position: 1, + presetName: 'first_preset', + run: makeRun({ runId: 'run-1' }), + }), + makeEntry({ + entryId: SECOND_ENTRY_ID, + position: 2, + presetName: 'second_preset', + state: 'QUEUED', + run: makeRun({ runId: 'run-2', status: 'QUEUED', completedAttacks: 0 }), + }), + ], + ...overrides, + } +} + +function setHookResult( + state: UseScenarioSuiteExecutionResult['state'], +): void { + mockUseScenarioSuiteExecution.mockReturnValue({ + state, + retry: mockRetry, + applyExecution: mockApplyExecution, + }) +} + +function renderPage() { + return render( + + + + } + /> + + + , + ) +} + +describe('ScenarioSuiteExecutionPage', () => { + beforeEach(() => { + jest.clearAllMocks() + setHookResult({ + loadStatus: 'ready', + execution: makeExecution(), + error: null, + stale: false, + }) + }) + + it('shows ordered child progress and links to ordinary scenario runs', () => { + renderPage() + + expect(screen.getByRole('heading', { name: 'nightly_suite' })).toBeInTheDocument() + const rows = screen.getAllByTestId(/^suite-execution-entry-/) + expect(rows).toHaveLength(2) + expect(within(rows[0]).getByText('first_preset')).toBeInTheDocument() + expect(within(rows[0]).getByText('2 of 10 attacks')).toBeInTheDocument() + expect(within(rows[0]).getByRole('link', { name: 'View run' })) + .toHaveAttribute('href', '/scanner-history/run-1') + expect(within(rows[1]).getByText('second_preset')).toBeInTheDocument() + expect(within(rows[1]).getAllByText('Queued')).toHaveLength(2) + }) + + it('surfaces entry submission failures without inventing a child run', () => { + setHookResult({ + loadStatus: 'ready', + execution: makeExecution({ + status: 'FAILED', + entries: [ + makeEntry({ + entryId: FIRST_ENTRY_ID, + position: 1, + presetName: 'first_preset', + state: 'SUBMISSION_FAILED', + error: 'Scenario could not initialize.', + }), + ], + }), + error: null, + stale: false, + }) + + renderPage() + + expect(screen.getByText('Scenario could not initialize.')).toBeInTheDocument() + expect(screen.getByText('Not created')).toBeInTheDocument() + expect(screen.getAllByText('Submission failed')).toHaveLength(2) + }) + + it('cancels pending submissions and child runs after confirmation', async () => { + const user = userEvent.setup() + const cancelled = makeExecution({ + status: 'CANCELLED', + entries: makeExecution().entries.map((entry) => ({ ...entry, state: 'CANCELLED' })), + }) + mockCancelExecution.mockResolvedValue(cancelled) + + renderPage() + + await user.click(screen.getByRole('button', { name: 'Cancel suite' })) + const dialog = await screen.findByRole('dialog') + await user.click(within(dialog).getByTestId('confirm-cancel-suite')) + + await waitFor(() => expect(mockCancelExecution).toHaveBeenCalledWith('execution-1')) + expect(mockApplyExecution).toHaveBeenCalledWith(cancelled) + }) + + it('keeps the cancellation dialog open when cancellation fails', async () => { + const user = userEvent.setup() + mockCancelExecution.mockRejectedValue(new Error('Cancellation unavailable.')) + + renderPage() + + await user.click(screen.getByRole('button', { name: 'Cancel suite' })) + const dialog = await screen.findByRole('dialog') + await user.click(within(dialog).getByTestId('confirm-cancel-suite')) + + expect(await within(dialog).findByText('Cancellation unavailable.')).toBeInTheDocument() + expect(mockApplyExecution).not.toHaveBeenCalled() + }) + + it('explains when transient suite progress is no longer available', () => { + setHookResult({ + loadStatus: 'not-found', + execution: null, + error: 'Execution is no longer available.', + stale: false, + }) + + renderPage() + + expect(screen.getByRole('heading', { name: 'Suite execution not found' })).toBeInTheDocument() + expect(screen.getByText(/successfully created scenario runs remain/i)).toBeInTheDocument() + }) +}) diff --git a/frontend/src/components/ScenarioSuites/ScenarioSuiteExecutionPage.tsx b/frontend/src/components/ScenarioSuites/ScenarioSuiteExecutionPage.tsx new file mode 100644 index 0000000000..2f8e0a5013 --- /dev/null +++ b/frontend/src/components/ScenarioSuites/ScenarioSuiteExecutionPage.tsx @@ -0,0 +1,431 @@ +import { type ReactElement, useState } from 'react' + +import { + Badge, + Button, + Dialog, + DialogActions, + DialogBody, + DialogContent, + DialogSurface, + DialogTitle, + MessageBar, + MessageBarBody, + mergeClasses, + Spinner, + Table, + TableBody, + TableCell, + TableHeader, + TableHeaderCell, + TableRow, + Text, +} from '@fluentui/react-components' +import { + ArrowLeftRegular, + ArrowSyncRegular, + CheckmarkCircleRegular, + DismissCircleRegular, + ErrorCircleRegular, + StopRegular, +} from '@fluentui/react-icons' +import { Link, useParams } from 'react-router' + +import { formatTimestamp } from '@/components/AttackResults/attackAttemptFormatting' +import { + isTerminalSuiteExecutionStatus, + useScenarioSuiteExecution, +} from '@/hooks/useScenarioSuiteExecution' +import { scenarioSuitesApi } from '@/services/api' +import { toApiError } from '@/services/errors' +import type { + ScenarioSuiteExecutionEntryState, + ScenarioSuiteExecutionEntrySummary, + ScenarioSuiteExecutionStatus, +} from '@/types' +import { routerPathParamValue, scenarioRunRoutePath } from '@/utils/routeParams' + +import { useScenarioSuiteExecutionPageStyles } from './ScenarioSuiteExecutionPage.styles' +import { SUITES_ROUTE } from './suiteRoutes' + +const EXECUTION_BADGE_COLORS: Record< + ScenarioSuiteExecutionStatus, + 'informative' | 'brand' | 'success' | 'danger' | 'warning' +> = { + QUEUING: 'informative', + IN_PROGRESS: 'brand', + COMPLETED: 'success', + FAILED: 'danger', + CANCELLED: 'warning', +} + +const ENTRY_BADGE_COLORS: Record< + ScenarioSuiteExecutionEntryState, + 'informative' | 'brand' | 'success' | 'danger' | 'warning' +> = { + PENDING: 'informative', + SUBMITTING: 'informative', + SUBMISSION_FAILED: 'danger', + CREATED: 'informative', + QUEUED: 'informative', + IN_PROGRESS: 'brand', + COMPLETED: 'success', + FAILED: 'danger', + CANCELLED: 'warning', +} + +export default function ScenarioSuiteExecutionPage() { + const { executionId } = useParams<{ executionId: string }>() + const decodedExecutionId = routerPathParamValue(executionId) + return ( + + ) +} + +interface ScenarioSuiteExecutionPageContentProps { + readonly executionId: string +} + +function ScenarioSuiteExecutionPageContent({ + executionId, +}: ScenarioSuiteExecutionPageContentProps) { + const styles = useScenarioSuiteExecutionPageStyles() + const { state, retry, applyExecution } = useScenarioSuiteExecution(executionId) + const [cancelDialogOpen, setCancelDialogOpen] = useState(false) + const [cancelling, setCancelling] = useState(false) + const [cancelError, setCancelError] = useState(null) + + const handleCancel = async (): Promise => { + setCancelling(true) + setCancelError(null) + try { + const execution = await scenarioSuitesApi.cancelExecution(executionId) + applyExecution(execution) + setCancelDialogOpen(false) + } catch (error: unknown) { + setCancelError(toApiError(error).detail) + } finally { + setCancelling(false) + } + } + + if (state.loadStatus === 'loading' && !state.execution) { + return ( +
+
+ + Back to suites + +
+ +
+
+
+ ) + } + + if (state.loadStatus === 'not-found' && !state.execution) { + return ( +
+
+ + Back to suites + +
+ + Suite execution not found + {state.error} + + Suite progress is process-local, but successfully created scenario runs remain in scanner history. + + +
+
+
+ ) + } + + if (state.loadStatus === 'error' && !state.execution) { + return ( +
+
+ + Back to suites + +
+ + Unable to load suite execution + {state.error} + +
+
+
+ ) + } + + const execution = state.execution + if (!execution) { + return null + } + + const canCancel = !isTerminalSuiteExecutionStatus(execution.status) + const completedEntries = execution.entries.filter((entry) => entry.state === 'COMPLETED').length + const failedEntries = execution.entries.filter( + (entry) => entry.state === 'FAILED' || entry.state === 'SUBMISSION_FAILED', + ).length + const activeEntries = execution.entries.length - completedEntries - failedEntries + - execution.entries.filter((entry) => entry.state === 'CANCELLED').length + + return ( +
+
+ + Back to suites + + +
+
+
+ + {execution.suite_name} + + + {formatState(execution.status)} + +
+ + Suite execution ID: {execution.suite_execution_id} + +
+ {canCancel && ( +
+ +
+ )} +
+ + {state.stale && state.error && ( + + + Live progress could not be refreshed: {state.error} + + + )} + +
+
+ Suite origin + + {execution.suite_origin === 'registered' ? 'Defined in code' : 'Stored'} + +
+
+ Created + {formatTimestamp(execution.created_at)} +
+
+ Last updated + {formatTimestamp(execution.updated_at)} +
+
+ +
+
+ + Suite progress + + Each entry remains an ordinary scenario run. +
+
+ + + + +
+
+ + + + Order + Preset + Status + Scenario run + Progress + + + + {execution.entries.map((entry) => ( + + {entry.entry_position} + +
+ {entry.preset_name} + + {entry.preset_origin === 'registered' ? 'Defined in code' : 'Stored preset'} + + {entry.error && ( + + {entry.error} + + )} +
+ + + + {formatState(entry.state)} + + + + {entry.run ? ( + + View run + + ) : ( + Not created + )} + + {entryProgress(entry)} + + ))} + +
+
+ + {isTerminalSuiteExecutionStatus(execution.status) + ? `Suite ${formatState(execution.status)}` + : `${completedEntries} of ${execution.entries.length} entries completed`} + +
+
+ + { if (!data.open && !cancelling) setCancelDialogOpen(false) }} + > + + + Cancel {execution.suite_name}? + + + Pending submissions and running child scenarios will be cancelled where possible. + Completed child runs are unaffected. + + {cancelError && ( + + {cancelError} + + )} + + + + + + + + +
+ ) +} + +function SummaryItem({ label, value }: { readonly label: string; readonly value: number }) { + const styles = useScenarioSuiteExecutionPageStyles() + return ( +
+ {label} + {value} +
+ ) +} + +function entryProgress(entry: ScenarioSuiteExecutionEntrySummary): string { + if (!entry.run) { + return entry.state === 'SUBMISSION_FAILED' ? 'Submission failed' : 'Waiting for run' + } + if (entry.run.status === 'QUEUED') { + return entry.run.queue_position ? `Queue position ${entry.run.queue_position}` : 'Queued' + } + if (entry.run.total_attacks <= 0) { + return formatState(entry.run.status) + } + return `${entry.run.completed_attacks} of ${entry.run.total_attacks} attacks` +} + +function statusIcon( + status: ScenarioSuiteExecutionStatus | ScenarioSuiteExecutionEntryState, +): ReactElement | undefined { + if (status === 'COMPLETED') { + return + } + if (status === 'FAILED' || status === 'SUBMISSION_FAILED') { + return + } + if (status === 'CANCELLED') { + return + } + return undefined +} + +function formatState(status: string): string { + return status.toLowerCase().replace(/_/g, ' ').replace( + /^\w/, + (letter: string) => letter.toUpperCase(), + ) +} diff --git a/frontend/src/components/ScenarioSuites/ScenarioSuiteLibrary.styles.ts b/frontend/src/components/ScenarioSuites/ScenarioSuiteLibrary.styles.ts index f353335850..fc91449a4e 100644 --- a/frontend/src/components/ScenarioSuites/ScenarioSuiteLibrary.styles.ts +++ b/frontend/src/components/ScenarioSuites/ScenarioSuiteLibrary.styles.ts @@ -65,7 +65,7 @@ export const useScenarioSuiteLibraryStyles = makeStyles({ overflowX: 'auto', }, table: { - minWidth: '760px', + minWidth: '900px', backgroundColor: tokens.colorNeutralBackground1, }, nameCell: { diff --git a/frontend/src/components/ScenarioSuites/ScenarioSuiteLibrary.test.tsx b/frontend/src/components/ScenarioSuites/ScenarioSuiteLibrary.test.tsx index a05fa243f6..9145efe8c7 100644 --- a/frontend/src/components/ScenarioSuites/ScenarioSuiteLibrary.test.tsx +++ b/frontend/src/components/ScenarioSuites/ScenarioSuiteLibrary.test.tsx @@ -4,6 +4,7 @@ import { FluentProvider, webLightTheme } from '@fluentui/react-components' import { MemoryRouter } from 'react-router' import { scenarioSuitesApi } from '@/services/api' +import { makeTarget } from '@/test-utils/targetFixtures' import type { ScenarioSuiteResponse } from '@/types' import ScenarioSuiteLibrary from './ScenarioSuiteLibrary' @@ -12,11 +13,13 @@ jest.mock('@/services/api', () => ({ scenarioSuitesApi: { list: jest.fn(), remove: jest.fn(), + launch: jest.fn(), }, })) const mockList = scenarioSuitesApi.list as jest.Mock const mockRemove = scenarioSuitesApi.remove as jest.Mock +const mockLaunch = scenarioSuitesApi.launch as jest.Mock const mockNavigate = jest.fn() jest.mock('react-router', () => ({ @@ -26,6 +29,7 @@ jest.mock('react-router', () => ({ const FIRST_ENTRY_ID = '11111111-1111-4111-8111-111111111111' const SECOND_ENTRY_ID = '22222222-2222-4222-8222-222222222222' +const OBJECTIVE_TARGET = makeTarget({ target_registry_name: 'objective-target' }) function makeSuite(): ScenarioSuiteResponse { return { @@ -68,11 +72,18 @@ function makeRegisteredSuite(): ScenarioSuiteResponse { } } -function renderLibrary() { +function renderLibrary({ defaultsReady = true }: { defaultsReady?: boolean } = {}) { return render( - + , ) @@ -82,6 +93,7 @@ describe('ScenarioSuiteLibrary', () => { beforeEach(() => { jest.clearAllMocks() mockList.mockResolvedValue({ source: 'suites', items: [] }) + mockLaunch.mockResolvedValue({ suite_execution_id: 'execution-1' }) }) it('shows an empty state and configured storage source', async () => { @@ -91,7 +103,7 @@ describe('ScenarioSuiteLibrary', () => { expect(screen.getByText('Stored in suites')).toBeInTheDocument() }) - it('lists duplicate preset references, author, and description without launch controls', async () => { + it('lists duplicate preset references, author, description, and launch controls', async () => { mockList.mockResolvedValue({ source: 'suites', items: [makeSuite()] }) renderLibrary() @@ -100,7 +112,7 @@ describe('ScenarioSuiteLibrary', () => { expect(within(row).getByText('2 entries')).toBeInTheDocument() expect(within(row).getByText('Nightly checks.')).toBeInTheDocument() expect(within(row).getByText('Ada Lovelace')).toBeInTheDocument() - expect(screen.queryByRole('button', { name: /launch/i })).not.toBeInTheDocument() + expect(within(row).getByRole('button', { name: 'Launch' })).toBeEnabled() }) it('surfaces entry-scoped advisory issues', async () => { @@ -121,6 +133,7 @@ describe('ScenarioSuiteLibrary', () => { expect(screen.getByText( "Entry 2 (retired_probe): Preset 'retired_probe' is not available.", )).toBeInTheDocument() + expect(screen.getByTestId('launch-suite-nightly_suite')).toBeDisabled() }) it('requires confirmation before deleting and refreshes afterwards', async () => { @@ -149,6 +162,7 @@ describe('ScenarioSuiteLibrary', () => { const row = await screen.findByTestId('suite-row-nightly_suite') expect(within(row).getByText('Defined in code')).toBeInTheDocument() + expect(within(row).getByRole('button', { name: 'Launch' })).toBeEnabled() expect(within(row).getByRole('button', { name: 'View' })).toBeEnabled() expect(within(row).getByRole('button', { name: 'Delete' })).toBeDisabled() @@ -157,6 +171,60 @@ describe('ScenarioSuiteLibrary', () => { expect(mockRemove).not.toHaveBeenCalled() }) + it('opens the launch dialog for a runnable suite', async () => { + const user = userEvent.setup() + mockList.mockResolvedValue({ source: 'suites', items: [makeSuite()] }) + + renderLibrary() + + await user.click(await screen.findByTestId('launch-suite-nightly_suite')) + + expect(await screen.findByTestId('confirm-launch-suite')).toBeInTheDocument() + expect(screen.getByTestId('launch-suite-entries')).toHaveTextContent( + 'nightly_probenightly_probe', + ) + }) + + it('refreshes the suite after a launch preflight conflict', async () => { + const user = userEvent.setup() + mockList.mockResolvedValue({ source: 'suites', items: [makeSuite()] }) + mockLaunch.mockRejectedValue({ + isAxiosError: true, + response: { status: 409, data: { detail: 'Suite changed.' } }, + }) + + renderLibrary() + + await user.click(await screen.findByTestId('launch-suite-nightly_suite')) + await user.click(screen.getByTestId('confirm-launch-suite')) + + expect(await screen.findByText( + "Suite 'nightly_suite' or one of its presets changed while you were launching it. " + + 'Review it and launch again.', + )).toBeInTheDocument() + await waitFor(() => expect(mockList).toHaveBeenCalledTimes(2)) + }) + + it('blocks launch while the global launch context is unresolved', async () => { + mockList.mockResolvedValue({ source: 'suites', items: [makeSuite()] }) + + renderLibrary({ defaultsReady: false }) + + expect(await screen.findByTestId('launch-suite-nightly_suite')).toBeDisabled() + }) + + it('does not launch an empty suite', async () => { + const suite = makeSuite() + suite.suite.entries = [] + suite.entry_statuses = [] + mockList.mockResolvedValue({ source: 'suites', items: [suite] }) + + renderLibrary() + + expect(await screen.findByText('Empty suite')).toBeInTheDocument() + expect(screen.getByTestId('launch-suite-nightly_suite')).toBeDisabled() + }) + it('reports load failures and retries', async () => { const user = userEvent.setup() mockList.mockRejectedValueOnce(new Error('suite storage unavailable')) diff --git a/frontend/src/components/ScenarioSuites/ScenarioSuiteLibrary.tsx b/frontend/src/components/ScenarioSuites/ScenarioSuiteLibrary.tsx index fadfa3959f..75d64c67d8 100644 --- a/frontend/src/components/ScenarioSuites/ScenarioSuiteLibrary.tsx +++ b/frontend/src/components/ScenarioSuites/ScenarioSuiteLibrary.tsx @@ -20,18 +20,36 @@ import { TableRow, Text, } from '@fluentui/react-components' -import { AddRegular, ArrowSyncRegular, DeleteRegular, EditRegular } from '@fluentui/react-icons' +import { + AddRegular, + ArrowSyncRegular, + DeleteRegular, + EditRegular, + PlayRegular, +} from '@fluentui/react-icons' import { Link, useNavigate } from 'react-router' import { scenarioSuitesApi } from '@/services/api' import { toApiError } from '@/services/errors' -import type { ScenarioSuiteResponse } from '@/types' +import type { ScenarioSuiteResponse, TargetInstance } from '@/types' +import LaunchSuiteDialog from './LaunchSuiteDialog' import { useScenarioSuiteLibraryStyles } from './ScenarioSuiteLibrary.styles' import { NEW_SUITE_ROUTE, suiteEditorRoutePath } from './suiteRoutes' const DEFINED_IN_CODE_HINT = 'Defined in code, so it cannot be edited or deleted here.' +interface ScenarioSuiteLibraryProps { + targets: TargetInstance[] + defaultObjectiveTarget: TargetInstance | null + defaultAdversarialTarget: TargetInstance | null + labels: Record + /** The saved Operation that new runs join, or null for none. */ + operationId: string | null + /** False while server defaults or the global Operation selection are unresolved. */ + defaultsReady: boolean +} + function issueCount(item: ScenarioSuiteResponse): number { return item.entry_statuses.reduce( (count, status) => count + status.issues.length, @@ -39,7 +57,14 @@ function issueCount(item: ScenarioSuiteResponse): number { ) } -export default function ScenarioSuiteLibrary() { +export default function ScenarioSuiteLibrary({ + targets, + defaultObjectiveTarget, + defaultAdversarialTarget, + labels, + operationId, + defaultsReady, +}: ScenarioSuiteLibraryProps) { const styles = useScenarioSuiteLibraryStyles() const navigate = useNavigate() const [items, setItems] = useState([]) @@ -47,6 +72,7 @@ export default function ScenarioSuiteLibrary() { const [loading, setLoading] = useState(true) const [error, setError] = useState(null) const [actionError, setActionError] = useState(null) + const [launching, setLaunching] = useState(null) const [pendingDelete, setPendingDelete] = useState(null) const [deleting, setDeleting] = useState(false) const [generation, setGeneration] = useState(0) @@ -190,6 +216,7 @@ export default function ScenarioSuiteLibrary() { {items.map((item) => { const issues = issueCount(item) + const runnable = item.suite.entries.length > 0 && issues === 0 const editable = item.origin === 'stored' return ( @@ -239,7 +266,9 @@ export default function ScenarioSuiteLibrary() { - {issues === 0 + {item.suite.entries.length === 0 + ? Empty suite + : issues === 0 ? No advisory issues : {issues} advisory issue{issues === 1 ? '' : 's'}} @@ -250,6 +279,16 @@ export default function ScenarioSuiteLibrary() {
+ + + + ) +} + +describe('useScenarioSuiteExecution', () => { + beforeEach(() => { + jest.useFakeTimers() + jest.clearAllMocks() + }) + + afterEach(() => { + jest.useRealTimers() + }) + + it('polls until the execution reaches a terminal state', async () => { + mockGetExecution + .mockResolvedValueOnce(makeExecution('IN_PROGRESS')) + .mockResolvedValueOnce(makeExecution('COMPLETED')) + + render() + + await waitFor(() => expect(screen.getByTestId('execution-status')).toHaveTextContent('IN_PROGRESS')) + await act(async () => { + await jest.advanceTimersByTimeAsync(SCENARIO_SUITE_EXECUTION_POLL_INTERVAL_MS) + }) + expect(screen.getByTestId('execution-status')).toHaveTextContent('COMPLETED') + + await act(async () => { + await jest.advanceTimersByTimeAsync(SCENARIO_SUITE_EXECUTION_POLL_INTERVAL_MS * 2) + }) + expect(mockGetExecution).toHaveBeenCalledTimes(2) + }) + + it('retries one transient polling failure before surfacing an error', async () => { + mockGetExecution + .mockRejectedValueOnce(new Error('temporary outage')) + .mockResolvedValueOnce(makeExecution('IN_PROGRESS')) + + render() + + await waitFor(() => expect(mockGetExecution).toHaveBeenCalledTimes(1)) + expect(screen.getByTestId('error')).toHaveTextContent('') + + await act(async () => { + await jest.advanceTimersByTimeAsync(SCENARIO_SUITE_EXECUTION_POLL_INTERVAL_MS) + }) + + expect(screen.getByTestId('execution-status')).toHaveTextContent('IN_PROGRESS') + expect(screen.getByTestId('load-status')).toHaveTextContent('ready') + }) + + it('reports an expired process-local execution as not found', async () => { + mockGetExecution.mockRejectedValue({ + isAxiosError: true, + response: { status: 404, data: { detail: 'Execution is no longer available.' } }, + }) + + render() + + await waitFor(() => expect(screen.getByTestId('load-status')).toHaveTextContent('not-found')) + expect(screen.getByTestId('error')).toHaveTextContent('Execution is no longer available.') + }) + + it('keeps the last good execution visible when a later poll fails', async () => { + mockGetExecution + .mockResolvedValueOnce(makeExecution('IN_PROGRESS')) + .mockRejectedValue(new Error('refresh failed')) + + render() + + await waitFor(() => expect(screen.getByTestId('execution-status')).toHaveTextContent('IN_PROGRESS')) + await act(async () => { + await jest.advanceTimersByTimeAsync(SCENARIO_SUITE_EXECUTION_POLL_INTERVAL_MS * 2) + }) + + expect(screen.getByTestId('execution-status')).toHaveTextContent('IN_PROGRESS') + expect(screen.getByTestId('stale')).toHaveTextContent('true') + expect(screen.getByTestId('error')).toHaveTextContent('refresh failed') + }) + + it('stops polling when cancellation applies a terminal response', async () => { + const user = userEvent.setup({ advanceTimers: jest.advanceTimersByTime }) + mockGetExecution.mockResolvedValue(makeExecution('IN_PROGRESS')) + + render() + + await waitFor(() => expect(screen.getByTestId('execution-status')).toHaveTextContent('IN_PROGRESS')) + await user.click(screen.getByRole('button', { name: 'Apply cancelled' })) + expect(screen.getByTestId('execution-status')).toHaveTextContent('CANCELLED') + + await act(async () => { + await jest.advanceTimersByTimeAsync(SCENARIO_SUITE_EXECUTION_POLL_INTERVAL_MS * 2) + }) + expect(mockGetExecution).toHaveBeenCalledTimes(1) + }) +}) diff --git a/frontend/src/hooks/useScenarioSuiteExecution.ts b/frontend/src/hooks/useScenarioSuiteExecution.ts new file mode 100644 index 0000000000..4880878e94 --- /dev/null +++ b/frontend/src/hooks/useScenarioSuiteExecution.ts @@ -0,0 +1,165 @@ +import { useCallback, useEffect, useRef, useState } from 'react' + +import { scenarioSuitesApi } from '@/services/api' +import { toApiError } from '@/services/errors' +import type { + ScenarioSuiteExecutionStatus, + ScenarioSuiteExecutionSummary, +} from '@/types' + +export const SCENARIO_SUITE_EXECUTION_POLL_INTERVAL_MS = 2_500 + +export type ScenarioSuiteExecutionLoadStatus = 'loading' | 'ready' | 'not-found' | 'error' + +export interface ScenarioSuiteExecutionState { + readonly loadStatus: ScenarioSuiteExecutionLoadStatus + readonly execution: ScenarioSuiteExecutionSummary | null + readonly error: string | null + readonly stale: boolean +} + +export interface UseScenarioSuiteExecutionResult { + readonly state: ScenarioSuiteExecutionState + readonly retry: () => void + readonly applyExecution: (execution: ScenarioSuiteExecutionSummary) => void +} + +const TERMINAL_STATUSES: ReadonlySet = new Set([ + 'COMPLETED', + 'FAILED', + 'CANCELLED', +]) + +export function isTerminalSuiteExecutionStatus(status: ScenarioSuiteExecutionStatus): boolean { + return TERMINAL_STATUSES.has(status) +} + +export function useScenarioSuiteExecution( + executionId: string, +): UseScenarioSuiteExecutionResult { + const executionIdValid = executionId.trim() !== '' + const [state, setState] = useState(() => ( + executionIdValid + ? { + loadStatus: 'loading', + execution: null, + error: null, + stale: false, + } + : { + loadStatus: 'not-found', + execution: null, + error: 'The scenario suite execution ID in this URL is missing or invalid.', + stale: false, + } + )) + const [retryEpoch, setRetryEpoch] = useState(0) + const timerRef = useRef | null>(null) + const controllerRef = useRef(null) + const latestExecutionRef = useRef(null) + const pollingStoppedRef = useRef(false) + const retriedRef = useRef(false) + + useEffect(() => { + let active = true + pollingStoppedRef.current = false + retriedRef.current = false + + const clearPollTimer = (): void => { + if (timerRef.current !== null) { + clearTimeout(timerRef.current) + timerRef.current = null + } + } + + const schedulePoll = (): void => { + clearPollTimer() + timerRef.current = setTimeout(() => { + timerRef.current = null + void fetchExecutionAsync() + }, SCENARIO_SUITE_EXECUTION_POLL_INTERVAL_MS) + } + + const fetchExecutionAsync = async (): Promise => { + if (!active || pollingStoppedRef.current) { + return + } + const controller = new AbortController() + controllerRef.current = controller + try { + const execution = await scenarioSuitesApi.getExecution(executionId, controller.signal) + if (!active || pollingStoppedRef.current) { + return + } + latestExecutionRef.current = execution + retriedRef.current = false + setState({ loadStatus: 'ready', execution, error: null, stale: false }) + if (isTerminalSuiteExecutionStatus(execution.status)) { + pollingStoppedRef.current = true + return + } + schedulePoll() + } catch (error: unknown) { + if (!active || controller.signal.aborted) { + return + } + const apiError = toApiError(error) + if (apiError.status !== 404 && !retriedRef.current) { + retriedRef.current = true + schedulePoll() + return + } + const previous = latestExecutionRef.current + setState({ + loadStatus: apiError.status === 404 && previous === null + ? 'not-found' + : previous === null + ? 'error' + : 'ready', + execution: previous, + error: apiError.detail, + stale: previous !== null, + }) + } + } + + if (executionIdValid) { + void fetchExecutionAsync() + } + + return () => { + active = false + clearPollTimer() + controllerRef.current?.abort() + controllerRef.current = null + } + }, [executionId, executionIdValid, retryEpoch]) + + const retry = useCallback((): void => { + pollingStoppedRef.current = false + retriedRef.current = false + setState((previous) => ({ + ...previous, + loadStatus: previous.execution ? 'ready' : 'loading', + error: null, + stale: false, + })) + setRetryEpoch((epoch) => epoch + 1) + }, []) + + const applyExecution = useCallback((execution: ScenarioSuiteExecutionSummary): void => { + latestExecutionRef.current = execution + setState({ loadStatus: 'ready', execution, error: null, stale: false }) + if (isTerminalSuiteExecutionStatus(execution.status)) { + pollingStoppedRef.current = true + if (timerRef.current !== null) { + clearTimeout(timerRef.current) + timerRef.current = null + } + controllerRef.current?.abort() + controllerRef.current = null + } + }, []) + + return { state, retry, applyExecution } +} diff --git a/frontend/src/services/api.test.ts b/frontend/src/services/api.test.ts index 8becaa3594..b7d5a0e384 100644 --- a/frontend/src/services/api.test.ts +++ b/frontend/src/services/api.test.ts @@ -135,12 +135,23 @@ describe("api service", () => { }); describe("scenarioSuitesApi", () => { - it("supports collection and encoded item CRUD with optimistic concurrency", async () => { + it("supports definition CRUD and transient execution APIs with encoded identifiers", async () => { const suite = { name: "nightly_suite", schema_version: 1, entries: [] }; const response = { suite, origin: "stored", version: "v1", entry_statuses: [] }; + const launch = { + expected_version: "v1", + entries: [], + target_name: "objective", + }; + const execution = { suite_execution_id: "execution/id" }; const controller = new AbortController(); - (apiClient.get as jest.Mock).mockResolvedValue({ data: response }); - (apiClient.post as jest.Mock).mockResolvedValue({ data: response }); + (apiClient.get as jest.Mock) + .mockResolvedValueOnce({ data: response }) + .mockResolvedValueOnce({ data: response }) + .mockResolvedValueOnce({ data: execution }); + (apiClient.post as jest.Mock) + .mockResolvedValueOnce({ data: response }) + .mockResolvedValue({ data: execution }); (apiClient.put as jest.Mock).mockResolvedValue({ data: response }); (apiClient.delete as jest.Mock).mockResolvedValue({ data: undefined }); @@ -157,6 +168,18 @@ describe("api service", () => { }); await scenarioSuitesApi.remove("a/b"); expect(apiClient.delete).toHaveBeenCalledWith("/scenario-suites/a%2Fb"); + await expect(scenarioSuitesApi.launch("a/b", launch)).resolves.toEqual(execution); + expect(apiClient.post).toHaveBeenCalledWith("/scenario-suites/a%2Fb/launch", launch); + await expect(scenarioSuitesApi.getExecution("execution/id", controller.signal)).resolves.toEqual(execution); + expect(apiClient.get).toHaveBeenCalledWith("/scenario-suites/executions/execution%2Fid", { + signal: controller.signal, + }); + await expect(scenarioSuitesApi.cancelExecution("execution/id", controller.signal)).resolves.toEqual(execution); + expect(apiClient.post).toHaveBeenCalledWith( + "/scenario-suites/executions/execution%2Fid/cancel", + undefined, + { signal: controller.signal }, + ); }); }); diff --git a/frontend/src/services/api.ts b/frontend/src/services/api.ts index 87ab32d7ac..20b9a01de2 100644 --- a/frontend/src/services/api.ts +++ b/frontend/src/services/api.ts @@ -52,6 +52,8 @@ import type { ScenarioPresetListResponse, ScenarioPresetResponse, ScenarioSuite, + LaunchScenarioSuiteRequest, + ScenarioSuiteExecutionSummary, ScenarioSuiteListResponse, ScenarioSuiteResponse, UpdateScenarioPresetRequest, @@ -829,6 +831,40 @@ export const scenarioSuitesApi = { remove: async (name: string): Promise => { await apiClient.delete(`/scenario-suites/${encodeURIComponent(name)}`) }, + + launch: async ( + name: string, + request: LaunchScenarioSuiteRequest, + ): Promise => { + const response = await apiClient.post( + `/scenario-suites/${encodeURIComponent(name)}/launch`, + request, + ) + return response.data + }, + + getExecution: async ( + executionId: string, + signal?: AbortSignal, + ): Promise => { + const response = await apiClient.get( + `/scenario-suites/executions/${encodeURIComponent(executionId)}`, + { signal }, + ) + return response.data + }, + + cancelExecution: async ( + executionId: string, + signal?: AbortSignal, + ): Promise => { + const response = await apiClient.post( + `/scenario-suites/executions/${encodeURIComponent(executionId)}/cancel`, + undefined, + { signal }, + ) + return response.data + }, } const MAX_SEED_SEARCH_LENGTH = 1_000 diff --git a/frontend/src/types/index.ts b/frontend/src/types/index.ts index bfb28f893b..f6431d4524 100644 --- a/frontend/src/types/index.ts +++ b/frontend/src/types/index.ts @@ -1449,6 +1449,71 @@ export interface UpdateScenarioSuiteRequest { expected_version: string } +/** The preset state the operator reviewed for one suite entry. */ +export interface ScenarioSuiteLaunchEntrySnapshot { + entry_id: string + preset_name: string + origin: 'stored' | 'registered' | null + preset_version: string | null +} + +/** Launch-owned fields and optimistic preconditions for one suite execution. */ +export interface LaunchScenarioSuiteRequest { + expected_version: string | null + entries: ScenarioSuiteLaunchEntrySnapshot[] + target_name: string + adversarial_target_name?: string | null + initializers?: string[] | null + initializer_args?: Record> | null + max_concurrency?: number | null + max_retries?: number | null + operation_id?: string | null + labels?: Record | null +} + +export type ScenarioSuiteExecutionStatus = + | 'QUEUING' + | 'IN_PROGRESS' + | 'COMPLETED' + | 'FAILED' + | 'CANCELLED' + +export type ScenarioSuiteExecutionEntryState = + | 'PENDING' + | 'SUBMITTING' + | 'SUBMISSION_FAILED' + | 'CREATED' + | 'QUEUED' + | 'IN_PROGRESS' + | 'COMPLETED' + | 'FAILED' + | 'CANCELLED' + +/** Current process-local state for one ordered suite entry. */ +export interface ScenarioSuiteExecutionEntrySummary { + entry_id: string + entry_position: number + preset_name: string + preset_origin: 'stored' | 'registered' + preset_version: string | null + state: ScenarioSuiteExecutionEntryState + run: ScenarioRunSummary | null + error: string | null + error_type: string | null +} + +/** Transient progress and grouped child runs for one suite launch. */ +export interface ScenarioSuiteExecutionSummary { + suite_execution_id: string + suite_name: string + suite_origin: 'stored' | 'registered' + suite_version: string | null + status: ScenarioSuiteExecutionStatus + created_at: string + updated_at: string + entries: ScenarioSuiteExecutionEntrySummary[] +} + /** The launch-owned fields a preset deliberately omits. */ export interface ResolveScenarioPresetRequest { expected_version?: string | null diff --git a/pyrit/backend/models/__init__.py b/pyrit/backend/models/__init__.py index 9b573fdc47..711e4270b9 100644 --- a/pyrit/backend/models/__init__.py +++ b/pyrit/backend/models/__init__.py @@ -71,7 +71,13 @@ UpdateScenarioPresetRequest, ) from pyrit.backend.models.scenario_suites import ( + LaunchScenarioSuiteRequest, ScenarioSuiteEntryStatus, + ScenarioSuiteExecutionEntryState, + ScenarioSuiteExecutionEntrySummary, + ScenarioSuiteExecutionStatus, + ScenarioSuiteExecutionSummary, + ScenarioSuiteLaunchEntrySnapshot, ScenarioSuiteListResponse, ScenarioSuiteResponse, UpdateScenarioSuiteRequest, @@ -132,6 +138,12 @@ "ScenarioPresetListResponse": "pyrit.backend.models.scenario_presets", "ScenarioPresetResponse": "pyrit.backend.models.scenario_presets", "UpdateScenarioPresetRequest": "pyrit.backend.models.scenario_presets", + "LaunchScenarioSuiteRequest": "pyrit.backend.models.scenario_suites", + "ScenarioSuiteExecutionEntryState": "pyrit.backend.models.scenario_suites", + "ScenarioSuiteExecutionEntrySummary": "pyrit.backend.models.scenario_suites", + "ScenarioSuiteExecutionStatus": "pyrit.backend.models.scenario_suites", + "ScenarioSuiteExecutionSummary": "pyrit.backend.models.scenario_suites", + "ScenarioSuiteLaunchEntrySnapshot": "pyrit.backend.models.scenario_suites", "ScenarioSuiteEntryStatus": "pyrit.backend.models.scenario_suites", "ScenarioSuiteListResponse": "pyrit.backend.models.scenario_suites", "ScenarioSuiteResponse": "pyrit.backend.models.scenario_suites", diff --git a/pyrit/backend/models/scenario_suites.py b/pyrit/backend/models/scenario_suites.py index e67073d4d9..f2de4d219b 100644 --- a/pyrit/backend/models/scenario_suites.py +++ b/pyrit/backend/models/scenario_suites.py @@ -1,17 +1,27 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -"""REST envelopes for scenario suite definition endpoints.""" +"""REST envelopes for scenario suite definition and execution endpoints.""" -from typing import Literal +from datetime import datetime +from enum import Enum +from typing import Any, Literal from uuid import UUID -from pydantic import BaseModel, Field +from pydantic import BaseModel, ConfigDict, Field +from pyrit.models.catalog.scenario import ScenarioRunSummary from pyrit.models.catalog.scenario_preset import PresetIssue from pyrit.models.catalog.scenario_suite import ScenarioSuite +from pyrit.models.request_limits import MAX_ITEMS __all__ = [ + "LaunchScenarioSuiteRequest", + "ScenarioSuiteExecutionEntryState", + "ScenarioSuiteExecutionEntrySummary", + "ScenarioSuiteExecutionStatus", + "ScenarioSuiteExecutionSummary", + "ScenarioSuiteLaunchEntrySnapshot", "ScenarioSuiteEntryStatus", "ScenarioSuiteListResponse", "ScenarioSuiteResponse", @@ -74,3 +84,127 @@ class UpdateScenarioSuiteRequest(BaseModel): suite: ScenarioSuite = Field(..., description="The replacement suite") expected_version: str = Field(..., description="Version returned when the suite being edited was read") + + +class ScenarioSuiteLaunchEntrySnapshot(BaseModel): + """The preset state the operator reviewed for one suite entry.""" + + model_config = ConfigDict(extra="forbid") + + entry_id: UUID = Field(..., description="Stable suite-entry identity") + preset_name: str = Field(..., description="Preset name shown for this entry") + origin: Literal["stored", "registered"] | None = Field( + ..., + description="Reviewed preset origin, or None when the entry was unresolved", + ) + preset_version: str | None = Field( + ..., + description="Reviewed stored preset version, or None for registered or unresolved presets", + ) + + +class LaunchScenarioSuiteRequest(BaseModel): + """Launch-owned fields and optimistic preconditions for one suite execution.""" + + model_config = ConfigDict(extra="forbid") + + expected_version: str | None = Field( + None, + description="Reviewed stored suite version; None only for a suite registered in code", + ) + entries: list[ScenarioSuiteLaunchEntrySnapshot] = Field( + ..., + min_length=1, + max_length=MAX_ITEMS, + description="Reviewed preset snapshots in the suite's declared order", + ) + target_name: str = Field(..., min_length=1, description="Registered objective target used by every entry") + adversarial_target_name: str | None = Field( + None, + min_length=1, + description="Registered adversarial target used by entries that require one", + ) + initializers: list[str] | None = Field( + None, + max_length=MAX_ITEMS, + description="Initializer names applied to every child scenario run", + ) + initializer_args: dict[str, dict[str, Any]] | None = Field( + None, + max_length=MAX_ITEMS, + description="Per-initializer parameter overrides keyed by initializer name", + ) + max_concurrency: int | None = Field( + None, + ge=1, + le=100, + description="Maximum concurrency applied independently to every child run", + ) + max_retries: int | None = Field( + None, + ge=0, + le=20, + description="Maximum retries applied independently to every child run", + ) + operation_id: UUID | None = Field(None, description="Saved Operation attributed to every child result") + labels: dict[str, str] | None = Field(None, description="Labels attached to every child result") + + +class ScenarioSuiteExecutionStatus(str, Enum): + """Process-local lifecycle of one accepted suite execution.""" + + QUEUING = "QUEUING" + IN_PROGRESS = "IN_PROGRESS" + COMPLETED = "COMPLETED" + FAILED = "FAILED" + CANCELLED = "CANCELLED" + + +class ScenarioSuiteExecutionEntryState(str, Enum): + """Submission or scenario-run state for one suite entry.""" + + PENDING = "PENDING" + SUBMITTING = "SUBMITTING" + SUBMISSION_FAILED = "SUBMISSION_FAILED" + CREATED = "CREATED" + QUEUED = "QUEUED" + IN_PROGRESS = "IN_PROGRESS" + COMPLETED = "COMPLETED" + FAILED = "FAILED" + CANCELLED = "CANCELLED" + + +class ScenarioSuiteExecutionEntrySummary(BaseModel): + """Current process-local state for one ordered suite entry.""" + + entry_id: UUID = Field(..., description="Stable suite-entry identity") + entry_position: int = Field(..., ge=1, description="One-based position in the resolved suite") + preset_name: str = Field(..., description="Resolved preset name") + preset_origin: Literal["stored", "registered"] = Field(..., description="Resolved preset origin") + preset_version: str | None = Field( + None, + description="Resolved stored preset version, or None for a registered preset", + ) + state: ScenarioSuiteExecutionEntryState = Field(..., description="Submission or child-run state") + run: ScenarioRunSummary | None = Field(None, description="Current child run summary after successful submission") + error: str | None = Field(None, description="Submission failure when no child run could be created") + error_type: str | None = Field(None, description="Submission exception class name") + + +class ScenarioSuiteExecutionSummary(BaseModel): + """Transient progress and grouped child results for one suite launch.""" + + suite_execution_id: UUID = Field(..., description="Server-owned identity for this suite launch") + suite_name: str = Field(..., description="Resolved suite name") + suite_origin: Literal["stored", "registered"] = Field(..., description="Resolved suite origin") + suite_version: str | None = Field( + None, + description="Resolved stored suite version, or None for a suite registered in code", + ) + status: ScenarioSuiteExecutionStatus = Field(..., description="Aggregate process-local execution state") + created_at: datetime = Field(..., description="When the suite execution was accepted") + updated_at: datetime = Field(..., description="Latest submission or child-run update") + entries: list[ScenarioSuiteExecutionEntrySummary] = Field( + ..., + description="Entry states in the suite's declared order", + ) diff --git a/pyrit/backend/routes/scenario_suites.py b/pyrit/backend/routes/scenario_suites.py index 0f514d4a21..07eeb4551f 100644 --- a/pyrit/backend/routes/scenario_suites.py +++ b/pyrit/backend/routes/scenario_suites.py @@ -4,17 +4,26 @@ """Scenario suite definition API routes.""" from typing import Annotated +from uuid import UUID from azure.core.exceptions import AzureError -from fastapi import APIRouter, Depends, HTTPException, status +from fastapi import APIRouter, Depends, HTTPException, Request, status from pyrit.backend.middleware.auth import current_user_name, require_admin from pyrit.backend.models.common import ProblemDetail from pyrit.backend.models.scenario_suites import ( + LaunchScenarioSuiteRequest, + ScenarioSuiteExecutionSummary, ScenarioSuiteListResponse, ScenarioSuiteResponse, UpdateScenarioSuiteRequest, ) +from pyrit.backend.services.scenario_suite_execution_service import ( + ScenarioSuiteExecutionNotFoundError, + ScenarioSuiteLaunchConflictError, + ScenarioSuiteNotRunnableError, + get_scenario_suite_execution_service, +) from pyrit.backend.services.scenario_suite_service import get_scenario_suite_service from pyrit.models.catalog import ScenarioSuite from pyrit.registry import ScenarioSuiteConflictError, ScenarioSuiteNameRegisteredError @@ -81,6 +90,100 @@ async def list_scenario_suites() -> ScenarioSuiteListResponse: # pyrit-async-su raise _storage_unavailable() from exc +@router.get( + "/executions/{execution_id}", + response_model=ScenarioSuiteExecutionSummary, + responses={404: {"model": ProblemDetail, "description": "Transient suite execution not found"}}, +) +async def get_scenario_suite_execution( # pyrit-async-suffix-exempt + execution_id: UUID, +) -> ScenarioSuiteExecutionSummary: + """ + Get process-local progress and grouped child runs for one suite execution. + + Returns: + ScenarioSuiteExecutionSummary: Current transient suite progress. + """ + execution = await get_scenario_suite_execution_service().get_execution_async(execution_id=execution_id) + if execution is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Scenario suite execution '{execution_id}' is no longer available", + ) + return execution + + +@router.post( + "/executions/{execution_id}/cancel", + response_model=ScenarioSuiteExecutionSummary, + responses={404: {"model": ProblemDetail, "description": "Transient suite execution not found"}}, +) +async def cancel_scenario_suite_execution( # pyrit-async-suffix-exempt + execution_id: UUID, +) -> ScenarioSuiteExecutionSummary: + """ + Best-effort cancel pending submissions and child scenario runs. + + Returns: + ScenarioSuiteExecutionSummary: Updated transient suite progress. + """ + execution = await get_scenario_suite_execution_service().cancel_execution_async(execution_id=execution_id) + if execution is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Scenario suite execution '{execution_id}' is no longer available", + ) + return execution + + +@router.post( + "/{name}/launch", + response_model=ScenarioSuiteExecutionSummary, + status_code=status.HTTP_202_ACCEPTED, + responses={ + 400: {"model": ProblemDetail, "description": "The suite has unresolved or invalid entries"}, + 404: {"model": ProblemDetail, "description": "Scenario suite not found"}, + 409: {"model": ProblemDetail, "description": "The suite or a referenced preset changed"}, + 503: {"model": ProblemDetail, "description": "Runtime or definition storage unavailable"}, + }, +) +async def launch_scenario_suite( # pyrit-async-suffix-exempt + name: str, + body: LaunchScenarioSuiteRequest, + http_request: Request, +) -> ScenarioSuiteExecutionSummary: + """ + Accept one transient suite execution after registry-backed preflight. + + Returns: + ScenarioSuiteExecutionSummary: Server-owned execution identity and + initial ordered entry states. + """ + runtime = getattr(http_request.app.state, "runtime_lifecycle", None) + if runtime is not None and runtime.state != "ready": + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail="Scenario suite launch was interrupted by runtime reinitialization.", + ) + try: + return await get_scenario_suite_execution_service().launch_async(name=name, request=body) + except ScenarioSuiteExecutionNotFoundError as exc: + raise _not_found(name) from exc + except ScenarioSuiteLaunchConflictError as exc: + raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail=str(exc)) from None + except ScenarioSuiteConflictError as exc: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=f"Scenario suite '{name}' changed after it was reviewed; refresh it before launching.", + ) from exc + except ScenarioSuiteNotRunnableError as exc: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None + except AzureError as exc: + raise _storage_unavailable() from exc + except RuntimeError as exc: + raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail=str(exc)) from None + + @router.get( "/{name}", response_model=ScenarioSuiteResponse, diff --git a/pyrit/backend/services/__init__.py b/pyrit/backend/services/__init__.py index fbf17703bf..b1ec37bf87 100644 --- a/pyrit/backend/services/__init__.py +++ b/pyrit/backend/services/__init__.py @@ -20,6 +20,10 @@ from pyrit.backend.services.scenario_preset_service import ScenarioPresetService, get_scenario_preset_service from pyrit.backend.services.scenario_run_service import ScenarioRunService, get_scenario_run_service from pyrit.backend.services.scenario_service import ScenarioService, get_scenario_service + from pyrit.backend.services.scenario_suite_execution_service import ( + ScenarioSuiteExecutionService, + get_scenario_suite_execution_service, + ) from pyrit.backend.services.scenario_suite_service import ScenarioSuiteService, get_scenario_suite_service from pyrit.backend.services.target_service import TargetService, get_target_service @@ -38,6 +42,8 @@ "get_scenario_preset_service": "pyrit.backend.services.scenario_preset_service", "ScenarioSuiteService": "pyrit.backend.services.scenario_suite_service", "get_scenario_suite_service": "pyrit.backend.services.scenario_suite_service", + "ScenarioSuiteExecutionService": "pyrit.backend.services.scenario_suite_execution_service", + "get_scenario_suite_execution_service": "pyrit.backend.services.scenario_suite_execution_service", "ScenarioRunService": "pyrit.backend.services.scenario_run_service", "get_scenario_run_service": "pyrit.backend.services.scenario_run_service", "TargetService": "pyrit.backend.services.target_service", diff --git a/pyrit/backend/services/runtime_lifecycle.py b/pyrit/backend/services/runtime_lifecycle.py index 67e6fcf0fb..9767c56749 100644 --- a/pyrit/backend/services/runtime_lifecycle.py +++ b/pyrit/backend/services/runtime_lifecycle.py @@ -16,6 +16,7 @@ from pyrit.backend.services.configuration_file_service import ConfigurationFileService from pyrit.backend.services.environment_file_service import EnvironmentFileService from pyrit.backend.services.scenario_run_service import get_scenario_run_service, peek_scenario_run_service +from pyrit.backend.services.scenario_suite_execution_service import peek_scenario_suite_execution_service from pyrit.backend.services.service_lifecycle import ( close_services_async, has_active_analytics_work, @@ -243,8 +244,10 @@ async def _apply_async(self, *, version: str) -> None: def _has_active_work(self) -> bool: """Return whether any admitted or background runtime operation remains.""" service = peek_scenario_run_service() + suite_execution_service = peek_scenario_suite_execution_service() return bool( (service and service.has_active_work()) + or (suite_execution_service and suite_execution_service.has_active_work()) or self.operations or outstanding_estimates() or has_active_manual_sends() @@ -268,6 +271,9 @@ async def shutdown_async(self) -> None: async def _shutdown_runtime_async(self) -> None: """Finish retained requests without cancelling their offloaded writes.""" service = peek_scenario_run_service() + suite_execution_service = peek_scenario_suite_execution_service() + if suite_execution_service: + suite_execution_service.stop_admission() if service: service.stop_admission() pending = self.operations | self.management_operations @@ -276,6 +282,13 @@ async def _shutdown_runtime_async(self) -> None: results = await asyncio.gather(*pending, return_exceptions=True) errors = [result for result in results if isinstance(result, BaseException)] self.state = "stopping" + try: + suite_execution_service = peek_scenario_suite_execution_service() + if suite_execution_service: + await suite_execution_service.shutdown_async() + except (Exception, asyncio.CancelledError) as error: + error.add_note("Runtime shutdown phase: scenario suite execution shutdown.") + errors.append(error) try: service = peek_scenario_run_service() if service: diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index 2936be6d7a..4d158ccc2a 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -17,6 +17,7 @@ from bisect import bisect_right from collections import OrderedDict, deque from collections.abc import AsyncGenerator, Mapping, Sequence +from contextvars import ContextVar from dataclasses import dataclass, replace from datetime import UTC, datetime from typing import Any @@ -101,6 +102,18 @@ _SHUTDOWN_INTERRUPTION_REASON = "The backend process shut down before this scenario run completed." _USER_CANCELLATION_REASON = "Run was cancelled by user" _LAUNCH_REQUEST_METADATA_KEY = "scenario_launch_request_v1" +_RESERVED_INITIAL_METADATA_KEYS = frozenset( + { + _SCHEDULER_METADATA_KEY, + _LAUNCH_REQUEST_METADATA_KEY, + SCENARIO_RUN_PLAN_METADATA_KEY, + SCENARIO_RUN_STARTED_AT_METADATA_KEY, + } +) +_INITIAL_METADATA_CONTEXT: ContextVar[Mapping[str, Any] | None] = ContextVar( + "scenario_run_initial_metadata", + default=None, +) _LAUNCH_REQUEST_FIELDS = ( "scenario_name", "target_name", @@ -246,16 +259,28 @@ def stop_admission(self) -> None: """Reject new preparations and prevent queued runs from starting.""" self._stopping = True - async def start_run_async(self, *, request: RunScenarioRequest) -> ScenarioRunSummary: + async def start_run_async( + self, + *, + request: RunScenarioRequest, + initial_metadata: Mapping[str, Any] | None = None, + ) -> ScenarioRunSummary: """ Initialize and schedule a new run or an explicitly configured continuation. + Args: + request: Scenario launch configuration. + initial_metadata: Internal caller-owned scenario metadata persisted + atomically with a fresh result. Reserved run-service keys cannot + be supplied. + Returns: ScenarioRunSummary: Current scheduled run state. """ + metadata = self._validate_initial_metadata(initial_metadata=initial_metadata) async with self._reserve_resume_request_async(request.scenario_result_id), self._launch_lock: await self._validate_resume_admission_async(scenario_result_id=request.scenario_result_id) - return await self._start_run_locked_async(request=request) + return await self._start_run_locked_async(request=request, initial_metadata=metadata) async def resume_run_async(self, *, scenario_result_id: str) -> ScenarioRunSummary: """ @@ -270,7 +295,7 @@ async def resume_run_async(self, *, scenario_result_id: str) -> ScenarioRunSumma ) assert stored is not None request = await asyncio.to_thread(self._restore_launch_request, stored=stored) - return await self._start_run_locked_async(request=request) + return await self._start_run_locked_async(request=request, initial_metadata={}) @contextlib.asynccontextmanager async def _reserve_resume_request_async(self, scenario_result_id: str | None) -> AsyncGenerator[None, None]: @@ -396,7 +421,12 @@ def _restore_launch_request(self, *, stored: ScenarioResult) -> RunScenarioReque deep=True, ) - async def _start_run_locked_async(self, *, request: RunScenarioRequest) -> ScenarioRunSummary: + async def _start_run_locked_async( + self, + *, + request: RunScenarioRequest, + initial_metadata: Mapping[str, Any], + ) -> ScenarioRunSummary: """ Initialize and schedule a scenario run. @@ -407,6 +437,8 @@ async def _start_run_locked_async(self, *, request: RunScenarioRequest) -> Scena Args: request: The run request with scenario name, target, and options. + initial_metadata: Internal caller-owned metadata persisted with a + fresh scenario result. Returns: ScenarioRunSummary with a stable ID and current active or queued state. @@ -426,7 +458,11 @@ async def _start_run_locked_async(self, *, request: RunScenarioRequest) -> Scena resumed_from_cancelled = await self._is_run_cancelled_async(scenario_result_id=request.scenario_result_id) if request.scenario_result_id: self._preparing_run_ids.add(request.scenario_result_id) - prepare_task = asyncio.create_task(self._prepare_run_async(request=request)) + metadata_token = _INITIAL_METADATA_CONTEXT.set(initial_metadata) + try: + prepare_task = asyncio.create_task(self._prepare_run_async(request=request)) + finally: + _INITIAL_METADATA_CONTEXT.reset(metadata_token) self._preparations.add(prepare_task) prepared = await asyncio.shield(prepare_task) return await self._schedule_prepared_run_async( @@ -611,12 +647,19 @@ async def _validate_run_operation_async(self, *, request: RunScenarioRequest) -> labels = {key: value for key, value in (request.labels or {}).items() if key != "operation"} return request.model_copy(update={"labels": labels if request.labels is not None else None}) - async def _prepare_run_async(self, *, request: RunScenarioRequest) -> _PreparedRun: + async def _prepare_run_async( + self, + *, + request: RunScenarioRequest, + initial_metadata: Mapping[str, Any] | None = None, + ) -> _PreparedRun: """ Resolve and initialize the scenario for a run request. Args: request: The run request with scenario name, target, and options. + initial_metadata: Internal caller-owned metadata persisted with a + fresh scenario result. Returns: _PreparedRun: The initialized scenario and its scoped default. @@ -663,9 +706,32 @@ async def _prepare_run_async(self, *, request: RunScenarioRequest) -> _PreparedR memory_labels=request.labels, ) init_kwargs["operation_id"] = request.operation_id - scenario = await self._initialize_scenario_async(request=request, init_kwargs=init_kwargs) + scenario = await self._initialize_scenario_async( + request=request, + init_kwargs=init_kwargs, + initial_metadata=( + initial_metadata if initial_metadata is not None else (_INITIAL_METADATA_CONTEXT.get() or {}) + ), + ) return _PreparedRun(scenario=scenario, adversarial_target=adversarial_target) + @staticmethod + def _validate_initial_metadata(*, initial_metadata: Mapping[str, Any] | None) -> dict[str, Any]: + """ + Copy caller metadata after protecting run-service-owned fields. + + Returns: + dict[str, Any]: A detached mapping safe to merge into scenario metadata. + + Raises: + ValueError: If the caller attempts to replace a reserved metadata key. + """ + metadata = dict(initial_metadata or {}) + reserved = sorted(_RESERVED_INITIAL_METADATA_KEYS.intersection(metadata)) + if reserved: + raise ValueError(f"Initial scenario metadata cannot set reserved keys: {', '.join(reserved)}") + return metadata + async def get_run_async(self, *, scenario_result_id: str) -> ScenarioRunSummary | None: """ Get the current status of a scenario run by querying the database. @@ -1249,7 +1315,13 @@ async def _run_initializers_async(self, *, request: RunScenarioRequest) -> None: raise ValueError(f"Initializer not found: {e}") from None await instance.initialize_async() - async def _initialize_scenario_async(self, *, request: RunScenarioRequest, init_kwargs: dict[str, Any]) -> Scenario: + async def _initialize_scenario_async( + self, + *, + request: RunScenarioRequest, + init_kwargs: dict[str, Any], + initial_metadata: Mapping[str, Any] | None = None, + ) -> Scenario: """ Build and initialize the scenario via the registry. @@ -1265,6 +1337,8 @@ async def _initialize_scenario_async(self, *, request: RunScenarioRequest, init_ scenario_result_id). init_kwargs: The resolved common parameters to pass to scenario.initialize_async. + initial_metadata: Internal caller-owned metadata persisted with a + fresh scenario result. Returns: The fully initialized Scenario instance ready for run_async. @@ -1331,6 +1405,7 @@ async def _initialize_scenario_async(self, *, request: RunScenarioRequest, init_ scenario_params=request.scenario_params or {}, scenario_result_id=request.scenario_result_id or None, initial_metadata={ + **(initial_metadata or {}), _SCHEDULER_METADATA_KEY: _SCHEDULER_METADATA_VALUE, _LAUNCH_REQUEST_METADATA_KEY: launch_request, }, diff --git a/pyrit/backend/services/scenario_suite_execution_service.py b/pyrit/backend/services/scenario_suite_execution_service.py new file mode 100644 index 0000000000..157e9e4e00 --- /dev/null +++ b/pyrit/backend/services/scenario_suite_execution_service.py @@ -0,0 +1,618 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Transient coordinator for launching registry-resolved scenario suites.""" + +from __future__ import annotations + +import asyncio +import logging +from collections import OrderedDict +from dataclasses import dataclass, field +from datetime import UTC, datetime +from typing import TYPE_CHECKING, Any, Literal +from uuid import UUID, uuid4 + +from pyrit.backend.models.scenario_presets import ResolveScenarioPresetRequest +from pyrit.backend.models.scenario_suites import ( + LaunchScenarioSuiteRequest, + ScenarioSuiteExecutionEntryState, + ScenarioSuiteExecutionEntrySummary, + ScenarioSuiteExecutionStatus, + ScenarioSuiteExecutionSummary, +) +from pyrit.backend.services.scenario_preset_service import ScenarioPresetService +from pyrit.backend.services.scenario_run_service import ScenarioRunService, get_scenario_run_service +from pyrit.models import ScenarioSuiteExecutionMetadata +from pyrit.registry import ScenarioPresetRegistry, ScenarioSuiteRegistry + +if TYPE_CHECKING: + from pyrit.models.catalog import RunScenarioRequest, ScenarioRunSummary, ScenarioSuiteRegistryEntry + +logger = logging.getLogger(__name__) + +_MAX_RETAINED_EXECUTIONS = 256 + + +class ScenarioSuiteExecutionNotFoundError(LookupError): + """A requested suite definition or transient execution does not exist.""" + + +class ScenarioSuiteLaunchConflictError(ValueError): + """A suite or referenced preset changed after the operator reviewed it.""" + + +class ScenarioSuiteNotRunnableError(ValueError): + """A suite cannot launch because one or more entries are unresolved.""" + + +@dataclass(frozen=True) +class _ResolvedLaunchEntry: + """One preflighted suite entry ready for scenario submission.""" + + entry_id: UUID + entry_position: int + preset_name: str + preset_origin: Literal["stored", "registered"] + preset_version: str | None + request: RunScenarioRequest + + def metadata(self, *, suite_execution_id: UUID, suite: ScenarioSuiteRegistryEntry) -> dict[str, Any]: + """Return the scenario metadata fields for this child run.""" + return ScenarioSuiteExecutionMetadata( + suite_execution_id=suite_execution_id, + suite_name=suite.suite.name, + suite_version=suite.version, + suite_entry_id=self.entry_id, + entry_position=self.entry_position, + preset_name=self.preset_name, + preset_version=self.preset_version, + ).to_metadata_fields() + + +@dataclass +class _ExecutionEntryState: + """Mutable process-local submission state for one resolved entry.""" + + resolved: _ResolvedLaunchEntry + state: ScenarioSuiteExecutionEntryState = ScenarioSuiteExecutionEntryState.PENDING + scenario_result_id: str | None = None + error: str | None = None + error_type: str | None = None + submission_task: asyncio.Task[ScenarioRunSummary] | None = None + + +@dataclass +class _ExecutionState: + """Mutable process-local state for one accepted suite execution.""" + + execution_id: UUID + suite: ScenarioSuiteRegistryEntry + entries: list[_ExecutionEntryState] + created_at: datetime + updated_at: datetime + cancel_requested: bool = False + task: asyncio.Task[None] | None = None + lock: asyncio.Lock = field(default_factory=asyncio.Lock) + + +@dataclass(frozen=True) +class _ExecutionEntrySnapshot: + """Detached entry state used while querying child runs.""" + + resolved: _ResolvedLaunchEntry + state: ScenarioSuiteExecutionEntryState + scenario_result_id: str | None + error: str | None + error_type: str | None + + +@dataclass(frozen=True) +class _ExecutionSnapshot: + """Detached execution state used while querying child runs.""" + + execution_id: UUID + suite: ScenarioSuiteRegistryEntry + entries: tuple[_ExecutionEntrySnapshot, ...] + created_at: datetime + updated_at: datetime + + +class ScenarioSuiteExecutionService: + """ + Preflight suites and coordinate ordinary scenario runs without a persisted parent. + + Definitions remain owned by the suite and preset registries. This service owns only + the transient identity and submission progress for one launch; each successful child + remains an ordinary persisted ``ScenarioResult`` managed by ``ScenarioRunService``. + """ + + def __init__( + self, + *, + suite_registry: ScenarioSuiteRegistry | None = None, + preset_registry: ScenarioPresetRegistry | None = None, + run_service: ScenarioRunService | None = None, + ) -> None: + """Initialize the coordinator with injectable registries and run service.""" + self._suite_registry = suite_registry + self._preset_registry = preset_registry + self._run_service = run_service + self._executions: OrderedDict[UUID, _ExecutionState] = OrderedDict() + self._tasks: set[asyncio.Task[None]] = set() + self._stopping = False + + def has_active_work(self) -> bool: + """Return whether accepted suite submissions are still running.""" + return any(not task.done() for task in self._tasks) + + def stop_admission(self) -> None: + """Reject new launches while allowing retained progress to remain readable.""" + self._stopping = True + + async def launch_async( + self, + *, + name: str, + request: LaunchScenarioSuiteRequest, + ) -> ScenarioSuiteExecutionSummary: + """ + Preflight one suite and begin ordered child-run submission. + + Returns: + ScenarioSuiteExecutionSummary: Accepted transient execution identity + and initial entry states. + """ + if self._stopping: + raise RuntimeError("Scenario suite launch scheduling is stopping.") + + suite, resolved_entries = await asyncio.to_thread( + self._resolve_launch, + name=name, + request=request, + ) + if self._stopping: + raise RuntimeError("Scenario suite launch scheduling is stopping.") + + now = datetime.now(UTC) + state = _ExecutionState( + execution_id=uuid4(), + suite=suite, + entries=[_ExecutionEntryState(resolved=entry) for entry in resolved_entries], + created_at=now, + updated_at=now, + ) + self._evict_completed_submissions() + self._executions[state.execution_id] = state + task = asyncio.create_task(self._submit_entries_async(state=state)) + state.task = task + self._tasks.add(task) + task.add_done_callback(self._submission_done) + return await self._build_summary_async(state=state) + + async def get_execution_async(self, *, execution_id: UUID) -> ScenarioSuiteExecutionSummary | None: + """ + Read current process-local progress for one suite execution. + + Returns: + ScenarioSuiteExecutionSummary | None: Current grouped progress, or + None when the transient execution is unavailable. + """ + state = self._executions.get(execution_id) + return None if state is None else await self._build_summary_async(state=state) + + async def cancel_execution_async(self, *, execution_id: UUID) -> ScenarioSuiteExecutionSummary | None: + """ + Best-effort cancel pending submissions and already-created child runs. + + Returns: + ScenarioSuiteExecutionSummary | None: Updated grouped progress, or + None when the transient execution is unavailable. + """ + state = self._executions.get(execution_id) + if state is None: + return None + + current = await self._build_summary_async(state=state) + if current.status in { + ScenarioSuiteExecutionStatus.COMPLETED, + ScenarioSuiteExecutionStatus.FAILED, + ScenarioSuiteExecutionStatus.CANCELLED, + }: + return current + + run_ids: list[str] = [] + async with state.lock: + state.cancel_requested = True + state.updated_at = datetime.now(UTC) + for entry in state.entries: + if entry.submission_task is not None and not entry.submission_task.done(): + entry.submission_task.cancel() + if entry.state == ScenarioSuiteExecutionEntryState.PENDING: + entry.state = ScenarioSuiteExecutionEntryState.CANCELLED + if entry.scenario_result_id is not None: + run_ids.append(entry.scenario_result_id) + + if run_ids: + await asyncio.gather( + *(self._get_run_service().cancel_run_async(scenario_result_id=run_id) for run_id in run_ids), + return_exceptions=True, + ) + return await self._build_summary_async(state=state) + + async def shutdown_async(self) -> None: + """Cancel retained submission tasks and wait for run-service cleanup.""" + self.stop_admission() + tasks = list(self._tasks) + for task in tasks: + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + + async def close_async(self) -> None: + """Close a drained coordinator.""" + if self.has_active_work(): + raise RuntimeError("Scenario suite execution work has not drained.") + + def _resolve_launch( + self, + *, + name: str, + request: LaunchScenarioSuiteRequest, + ) -> tuple[ScenarioSuiteRegistryEntry, list[_ResolvedLaunchEntry]]: + """ + Resolve and validate the complete launch before any scenario is submitted. + + Returns: + tuple[ScenarioSuiteRegistryEntry, list[_ResolvedLaunchEntry]]: The + exact suite and ordered, runnable child requests. + """ + suite = self._get_suite_registry().resolve_suite(name=name, expected_version=request.expected_version) + if suite is None: + raise ScenarioSuiteExecutionNotFoundError(f"Scenario suite '{name}' not found.") + if suite.origin == "stored" and request.expected_version is None: + raise ScenarioSuiteLaunchConflictError( + f"Scenario suite '{name}' must be refreshed before launch because no reviewed version was supplied." + ) + if len(suite.suite.entries) != len(request.entries): + raise ScenarioSuiteLaunchConflictError( + f"Scenario suite '{name}' changed after it was reviewed; refresh it before launching." + ) + + preset_registry = self._get_preset_registry() + launch_fields = ResolveScenarioPresetRequest( + expected_version=None, + target_name=request.target_name, + adversarial_target_name=request.adversarial_target_name, + initializers=request.initializers, + initializer_args=request.initializer_args, + max_concurrency=request.max_concurrency, + max_retries=request.max_retries, + operation_id=request.operation_id, + labels=request.labels, + ) + presets = { + preset_name: preset_registry.get_preset(name=preset_name) + for preset_name in dict.fromkeys(entry.preset_name for entry in suite.suite.entries) + } + + resolved: list[_ResolvedLaunchEntry] = [] + for position, (suite_entry, reviewed) in enumerate( + zip(suite.suite.entries, request.entries, strict=True), + start=1, + ): + if suite_entry.entry_id != reviewed.entry_id or suite_entry.preset_name != reviewed.preset_name: + raise ScenarioSuiteLaunchConflictError( + f"Scenario suite '{name}' changed after it was reviewed; refresh it before launching." + ) + + preset = presets[suite_entry.preset_name] + if preset is None: + raise ScenarioSuiteNotRunnableError( + f"Entry {position} references unavailable scenario preset '{suite_entry.preset_name}'." + ) + if preset.origin != reviewed.origin or preset.version != reviewed.preset_version: + raise ScenarioSuiteLaunchConflictError( + f"Scenario preset '{suite_entry.preset_name}' changed after the suite was reviewed; " + "refresh the suite before launching." + ) + + issues = preset_registry.check_preset(preset=preset.preset) + if issues: + details = "; ".join(issue.message for issue in issues) + raise ScenarioSuiteNotRunnableError( + f"Entry {position} scenario preset '{suite_entry.preset_name}' is not runnable: {details}" + ) + + resolved.append( + _ResolvedLaunchEntry( + entry_id=suite_entry.entry_id, + entry_position=position, + preset_name=suite_entry.preset_name, + preset_origin=preset.origin, + preset_version=preset.version, + request=ScenarioPresetService.resolve_run_request( + preset=preset.preset, + launch=launch_fields, + ), + ) + ) + return suite, resolved + + async def _submit_entries_async(self, *, state: _ExecutionState) -> None: + """Submit every resolved entry in declared order through the FIFO run service.""" + try: + async with state.lock: + if state.cancel_requested: + for entry in state.entries: + entry.state = ScenarioSuiteExecutionEntryState.CANCELLED + state.updated_at = datetime.now(UTC) + return + + for entry in state.entries: + entry.state = ScenarioSuiteExecutionEntryState.SUBMITTING + entry.submission_task = asyncio.create_task( + self._get_run_service().start_run_async( + request=entry.resolved.request, + initial_metadata=entry.resolved.metadata( + suite_execution_id=state.execution_id, + suite=state.suite, + ), + ) + ) + state.updated_at = datetime.now(UTC) + + for entry in state.entries: + task = entry.submission_task + if task is None: + continue + try: + result = await task + except asyncio.CancelledError: + async with state.lock: + entry.state = ScenarioSuiteExecutionEntryState.CANCELLED + state.updated_at = datetime.now(UTC) + continue + except Exception as exc: + logger.warning( + "Scenario suite execution %s entry %s could not be submitted.", + state.execution_id, + entry.resolved.entry_id, + exc_info=True, + ) + async with state.lock: + entry.state = ScenarioSuiteExecutionEntryState.SUBMISSION_FAILED + entry.error = str(exc) + entry.error_type = type(exc).__name__ + state.updated_at = datetime.now(UTC) + continue + + async with state.lock: + entry.scenario_result_id = result.scenario_result_id + entry.state = ScenarioSuiteExecutionEntryState(result.status.value) + should_cancel = state.cancel_requested + state.updated_at = datetime.now(UTC) + if should_cancel: + await self._get_run_service().cancel_run_async(scenario_result_id=result.scenario_result_id) + except asyncio.CancelledError: + async with state.lock: + tasks = [ + entry.submission_task + for entry in state.entries + if entry.submission_task is not None and not entry.submission_task.done() + ] + for task in tasks: + task.cancel() + for entry in state.entries: + if entry.state in { + ScenarioSuiteExecutionEntryState.PENDING, + ScenarioSuiteExecutionEntryState.SUBMITTING, + }: + entry.state = ScenarioSuiteExecutionEntryState.CANCELLED + state.updated_at = datetime.now(UTC) + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + raise + + async def _build_summary_async(self, *, state: _ExecutionState) -> ScenarioSuiteExecutionSummary: + """ + Build a coherent response while querying child runs outside the state lock. + + Returns: + ScenarioSuiteExecutionSummary: Current grouped execution progress. + """ + async with state.lock: + snapshot = _ExecutionSnapshot( + execution_id=state.execution_id, + suite=state.suite, + entries=tuple( + _ExecutionEntrySnapshot( + resolved=entry.resolved, + state=entry.state, + scenario_result_id=entry.scenario_result_id, + error=entry.error, + error_type=entry.error_type, + ) + for entry in state.entries + ), + created_at=state.created_at, + updated_at=state.updated_at, + ) + return await self._build_summary_from_snapshot_async(snapshot=snapshot) + + async def _build_summary_from_snapshot_async( + self, + *, + snapshot: _ExecutionSnapshot, + ) -> ScenarioSuiteExecutionSummary: + """ + Build a response from detached coordinator state and live child runs. + + Returns: + ScenarioSuiteExecutionSummary: Current grouped execution progress. + """ + run_entries = [entry for entry in snapshot.entries if entry.scenario_result_id is not None] + results = await asyncio.gather( + *( + self._get_run_service().get_run_async(scenario_result_id=entry.scenario_result_id or "") + for entry in run_entries + ), + return_exceptions=True, + ) + runs_by_entry: dict[UUID, ScenarioRunSummary | BaseException | None] = { + entry.resolved.entry_id: result for entry, result in zip(run_entries, results, strict=True) + } + + summaries: list[ScenarioSuiteExecutionEntrySummary] = [] + updated_at = snapshot.updated_at + for entry in snapshot.entries: + has_run = entry.resolved.entry_id in runs_by_entry + run_or_error = runs_by_entry.get(entry.resolved.entry_id) + run = None + error = entry.error + error_type = entry.error_type + entry_state = entry.state + if isinstance(run_or_error, BaseException): + entry_state = ScenarioSuiteExecutionEntryState.FAILED + error = str(run_or_error) + error_type = type(run_or_error).__name__ + elif run_or_error is not None: + run = run_or_error + entry_state = ScenarioSuiteExecutionEntryState(run.status.value) + updated_at = max(updated_at, run.updated_at) + elif has_run: + entry_state = ScenarioSuiteExecutionEntryState.FAILED + error = "The submitted child scenario run is no longer available." + error_type = "ScenarioRunNotFoundError" + + summaries.append( + ScenarioSuiteExecutionEntrySummary( + entry_id=entry.resolved.entry_id, + entry_position=entry.resolved.entry_position, + preset_name=entry.resolved.preset_name, + preset_origin=entry.resolved.preset_origin, + preset_version=entry.resolved.preset_version, + state=entry_state, + run=run, + error=error, + error_type=error_type, + ) + ) + + return ScenarioSuiteExecutionSummary( + suite_execution_id=snapshot.execution_id, + suite_name=snapshot.suite.suite.name, + suite_origin=snapshot.suite.origin, + suite_version=snapshot.suite.version, + status=self._execution_status(entries=summaries), + created_at=snapshot.created_at, + updated_at=updated_at, + entries=summaries, + ) + + @staticmethod + def _execution_status( + *, + entries: list[ScenarioSuiteExecutionEntrySummary], + ) -> ScenarioSuiteExecutionStatus: + """ + Reduce ordered entry states to one process-local execution status. + + Returns: + ScenarioSuiteExecutionStatus: Aggregate execution state. + """ + states = {entry.state for entry in entries} + if states.intersection( + { + ScenarioSuiteExecutionEntryState.PENDING, + ScenarioSuiteExecutionEntryState.SUBMITTING, + } + ): + return ScenarioSuiteExecutionStatus.QUEUING + if states.intersection( + { + ScenarioSuiteExecutionEntryState.CREATED, + ScenarioSuiteExecutionEntryState.QUEUED, + ScenarioSuiteExecutionEntryState.IN_PROGRESS, + } + ): + return ScenarioSuiteExecutionStatus.IN_PROGRESS + if ScenarioSuiteExecutionEntryState.CANCELLED in states: + return ScenarioSuiteExecutionStatus.CANCELLED + if states.intersection( + { + ScenarioSuiteExecutionEntryState.SUBMISSION_FAILED, + ScenarioSuiteExecutionEntryState.FAILED, + } + ): + return ScenarioSuiteExecutionStatus.FAILED + return ScenarioSuiteExecutionStatus.COMPLETED + + def _evict_completed_submissions(self) -> None: + """Bound transient coordinator state after submission tasks have finished.""" + while len(self._executions) >= _MAX_RETAINED_EXECUTIONS: + execution_id = next( + ( + execution_id + for execution_id, state in self._executions.items() + if state.task is not None and state.task.done() + ), + None, + ) + if execution_id is None: + return + self._executions.pop(execution_id, None) + + def _submission_done(self, task: asyncio.Task[None]) -> None: + """Retrieve unexpected task failures and release active-task ownership.""" + self._tasks.discard(task) + if task.cancelled(): + return + error = task.exception() + if error is not None: + logger.error( + "Scenario suite submission coordinator failed.", + exc_info=(type(error), error, error.__traceback__), + ) + + def _get_suite_registry(self) -> ScenarioSuiteRegistry: + """Return the injected suite registry or process singleton.""" + return ( + self._suite_registry if self._suite_registry is not None else ScenarioSuiteRegistry.get_registry_singleton() + ) + + def _get_preset_registry(self) -> ScenarioPresetRegistry: + """Return the injected preset registry or process singleton.""" + return ( + self._preset_registry + if self._preset_registry is not None + else ScenarioPresetRegistry.get_registry_singleton() + ) + + def _get_run_service(self) -> ScenarioRunService: + """Return the injected run service or process singleton.""" + return self._run_service if self._run_service is not None else get_scenario_run_service() + + +_service_instance: ScenarioSuiteExecutionService | None = None + + +def peek_scenario_suite_execution_service() -> ScenarioSuiteExecutionService | None: + """Return the existing coordinator without constructing one.""" + return _service_instance + + +async def reset_scenario_suite_execution_service_async() -> None: + """Close and discard the drained coordinator singleton.""" + global _service_instance + if _service_instance is not None: + await _service_instance.close_async() + _service_instance = None + + +def get_scenario_suite_execution_service() -> ScenarioSuiteExecutionService: + """Return the process-local suite execution coordinator.""" + global _service_instance + if _service_instance is None: + _service_instance = ScenarioSuiteExecutionService() + return _service_instance diff --git a/pyrit/backend/services/service_lifecycle.py b/pyrit/backend/services/service_lifecycle.py index 6e48e2c9fc..3ad17273da 100644 --- a/pyrit/backend/services/service_lifecycle.py +++ b/pyrit/backend/services/service_lifecycle.py @@ -14,6 +14,9 @@ from pyrit.backend.services.message_send_service import get_message_send_service from pyrit.backend.services.scenario_run_service import reset_scenario_run_service_async from pyrit.backend.services.scenario_service import get_scenario_service +from pyrit.backend.services.scenario_suite_execution_service import ( + reset_scenario_suite_execution_service_async, +) from pyrit.backend.services.scorer_service import get_scorer_service from pyrit.backend.services.target_service import get_target_service from pyrit.backend.services.technique_service import get_technique_service @@ -64,6 +67,7 @@ async def _close_cached_services_async() -> None: await get_scenario_service().close_async() if get_converter_service.cache_info().currsize: await get_converter_service().close_async() + await reset_scenario_suite_execution_service_async() await reset_scenario_run_service_async() finally: for factory in ( diff --git a/pyrit/models/__init__.py b/pyrit/models/__init__.py index cc6dcc462f..2db5d59b0e 100644 --- a/pyrit/models/__init__.py +++ b/pyrit/models/__init__.py @@ -195,6 +195,11 @@ ScenarioSeedGroupProgress, ScenarioTechniqueProgress, ) + from pyrit.models.scenario_suite_execution import ( + SCENARIO_SUITE_EXECUTION_METADATA_KEY, + SCENARIO_SUITE_EXECUTION_METADATA_VERSION, + ScenarioSuiteExecutionMetadata, + ) from pyrit.models.score import ( Acquisition, AnswerMatches, @@ -469,6 +474,9 @@ "SCENARIO_RUN_PLAN_METADATA_KEY": "pyrit.models.scenario_progress", "SCENARIO_RUN_PLAN_VERSION": "pyrit.models.scenario_progress", "SCENARIO_RUN_STARTED_AT_METADATA_KEY": "pyrit.models.scenario_progress", + "SCENARIO_SUITE_EXECUTION_METADATA_KEY": "pyrit.models.scenario_suite_execution", + "SCENARIO_SUITE_EXECUTION_METADATA_VERSION": "pyrit.models.scenario_suite_execution", + "ScenarioSuiteExecutionMetadata": "pyrit.models.scenario_suite_execution", "ScenarioAttackResultDelta": "pyrit.models.scenario_progress", "ScenarioAtomicGroupProgress": "pyrit.models.scenario_progress", "ScenarioAttackTechniqueDetails": "pyrit.models.scenario_progress", diff --git a/pyrit/models/scenario_suite_execution.py b/pyrit/models/scenario_suite_execution.py new file mode 100644 index 0000000000..9312a6dc5a --- /dev/null +++ b/pyrit/models/scenario_suite_execution.py @@ -0,0 +1,53 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Versioned metadata linking ordinary scenario results to one suite execution.""" + +from typing import Any, Literal, Self +from uuid import UUID + +from pydantic import BaseModel, ConfigDict, Field + +SCENARIO_SUITE_EXECUTION_METADATA_KEY = "scenario_suite_execution_v1" +SCENARIO_SUITE_EXECUTION_METADATA_VERSION = 1 + + +class ScenarioSuiteExecutionMetadata(BaseModel): + """Identity stamped on every scenario result successfully submitted by a suite launch.""" + + model_config = ConfigDict(extra="forbid") + + version: Literal[1] = SCENARIO_SUITE_EXECUTION_METADATA_VERSION + suite_execution_id: UUID = Field(..., description="Server-owned identity for one suite launch") + suite_name: str = Field(..., description="Resolved suite name") + suite_version: str | None = Field( + None, + description="Resolved stored suite version, or None for a suite registered in code", + ) + suite_entry_id: UUID = Field(..., description="Stable identity of the launched suite entry") + entry_position: int = Field(..., ge=1, description="One-based position in the resolved suite") + preset_name: str = Field(..., description="Resolved preset name") + preset_version: str | None = Field( + None, + description="Resolved stored preset version, or None for a preset registered in code", + ) + + def to_metadata_fields(self) -> dict[str, Any]: + """ + Return the namespaced scenario metadata entry owned by this model. + + Returns: + dict[str, Any]: JSON-compatible fields for ``ScenarioResult.metadata``. + """ + return {SCENARIO_SUITE_EXECUTION_METADATA_KEY: self.model_dump(mode="json")} + + @classmethod + def from_metadata(cls, metadata: dict[str, Any]) -> Self | None: + """ + Read suite execution metadata from a scenario result. + + Returns: + Self | None: The validated envelope, or None for a direct or legacy run. + """ + raw = metadata.get(SCENARIO_SUITE_EXECUTION_METADATA_KEY) + return None if raw is None else cls.model_validate(raw) diff --git a/tests/unit/backend/test_runtime_lifecycle.py b/tests/unit/backend/test_runtime_lifecycle.py index 9682c7f1d4..5bc4187327 100644 --- a/tests/unit/backend/test_runtime_lifecycle.py +++ b/tests/unit/backend/test_runtime_lifecycle.py @@ -315,6 +315,18 @@ async def test_active_work_rejects_apply_without_stopping_or_mutating(runtime: R assert not hasattr(service, "request_stop") or not service.request_stop.called +async def test_active_suite_submission_rejects_apply_without_mutating(runtime: RuntimeLifecycle) -> None: + """Test process-local suite coordination participates in the idle barrier.""" + service = MagicMock() + service.has_active_work.return_value = True + with patch.object(lifecycle_module, "peek_scenario_suite_execution_service", return_value=service): + await apply_async(runtime) + + assert runtime.outcome == "busy" + assert runtime.state == "ready" + lifecycle_module.close_services_async.assert_not_awaited() + + @pytest.mark.parametrize("count", [1, 3]) async def test_admitted_manual_send_blocks_reinitialization_without_an_http_request_async( runtime: RuntimeLifecycle, diff --git a/tests/unit/backend/test_scenario_run_service.py b/tests/unit/backend/test_scenario_run_service.py index f29c7a64e1..809275db9c 100644 --- a/tests/unit/backend/test_scenario_run_service.py +++ b/tests/unit/backend/test_scenario_run_service.py @@ -314,6 +314,47 @@ async def test_prepare_and_restore_preserve_common_input_support_async( await service.close_async() +async def test_prepare_persists_internal_initial_metadata_with_run_service_fields( + patch_central_database: MagicMock, +) -> None: + """Test suite attribution is created atomically with the scenario result.""" + registry = MagicMock(spec=ScenarioRegistry) + registry.get_class.return_value = Scam + registry.create_and_initialize_async = AsyncMock(return_value=MagicMock(spec=Scenario)) + service = ScenarioRunService() + suite_metadata = {"version": 1, "suite_execution_id": "execution-id"} + with ( + patch.object(ScenarioRegistry, "get_registry_singleton", return_value=registry), + patch.object(service, "_run_initializers_async", new_callable=AsyncMock), + patch.object(service._configuration_resolver, "resolve_target"), + ): + await service._prepare_run_async( + request=RunScenarioRequest(scenario_name="test.support", target_name="my_target"), + initial_metadata={"scenario_suite_execution_v1": suite_metadata}, + ) + + metadata = registry.create_and_initialize_async.call_args.kwargs["initial_metadata"] + assert metadata["scenario_suite_execution_v1"] == suite_metadata + assert metadata[_svc_mod._SCHEDULER_METADATA_KEY] == _svc_mod._SCHEDULER_METADATA_VALUE + assert _svc_mod._LAUNCH_REQUEST_METADATA_KEY in metadata + await service.close_async() + + +async def test_start_rejects_caller_metadata_that_overwrites_run_service_fields( + patch_central_database: MagicMock, +) -> None: + """Test internal launch callers cannot replace scheduler-owned metadata.""" + service = ScenarioRunService() + + with pytest.raises(ValueError, match="reserved keys"): + await service.start_run_async( + request=_make_request(), + initial_metadata={_svc_mod._LAUNCH_REQUEST_METADATA_KEY: {"forged": True}}, + ) + + await service.close_async() + + async def test_prepare_requires_target_for_scenarios_that_declare_it_async( patch_central_database: MagicMock, ) -> None: diff --git a/tests/unit/backend/test_scenario_suite_execution_service.py b/tests/unit/backend/test_scenario_suite_execution_service.py new file mode 100644 index 0000000000..a2d195c64d --- /dev/null +++ b/tests/unit/backend/test_scenario_suite_execution_service.py @@ -0,0 +1,472 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Tests for transient scenario suite execution coordination.""" + +import asyncio +from collections.abc import AsyncGenerator +from datetime import UTC, datetime +from pathlib import Path +from typing import Any, Literal +from unittest.mock import AsyncMock, MagicMock +from uuid import UUID + +import pytest + +from pyrit.backend.models.scenario_suites import ( + LaunchScenarioSuiteRequest, + ScenarioSuiteExecutionEntryState, + ScenarioSuiteExecutionStatus, + ScenarioSuiteLaunchEntrySnapshot, +) +from pyrit.backend.services.scenario_run_service import ScenarioRunService +from pyrit.backend.services.scenario_suite_execution_service import ( + ScenarioSuiteExecutionService, + ScenarioSuiteLaunchConflictError, + ScenarioSuiteNotRunnableError, +) +from pyrit.models import SCENARIO_SUITE_EXECUTION_METADATA_KEY, ScenarioRunState +from pyrit.models.catalog import ( + PresetEntry, + PresetIssue, + RunScenarioRequest, + ScenarioPreset, + ScenarioRunSummary, + ScenarioSuite, + ScenarioSuiteEntry, +) +from pyrit.registry import ScenarioPresetRegistry, ScenarioSuiteRegistry, ScenarioSuiteStorage + +ENTRY_ONE = UUID("11111111-1111-4111-8111-111111111111") +ENTRY_TWO = UUID("22222222-2222-4222-8222-222222222222") +OPERATION_ID = UUID("aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa") + + +def _suite(*, name: str = "weekly") -> ScenarioSuite: + """Build a two-entry suite with stable identities.""" + return ScenarioSuite( + name=name, + entries=[ + ScenarioSuiteEntry(entry_id=ENTRY_ONE, preset_name="first"), + ScenarioSuiteEntry(entry_id=ENTRY_TWO, preset_name="second"), + ], + ) + + +def _preset_entry( + *, + name: str, + scenario_name: str, + origin: Literal["stored", "registered"] = "stored", + version: str | None = None, +) -> PresetEntry: + """Build a registry entry for one runnable preset.""" + return PresetEntry( + preset=ScenarioPreset(name=name, scenario_name=scenario_name), + origin=origin, + version=version, + ) + + +def _run_summary( + *, + run_id: str, + scenario_name: str, + state: ScenarioRunState = ScenarioRunState.QUEUED, +) -> ScenarioRunSummary: + """Build a minimal child run summary.""" + now = datetime(2025, 1, 1, tzinfo=UTC) + return ScenarioRunSummary( + scenario_result_id=run_id, + scenario_name=scenario_name, + status=state, + created_at=now, + updated_at=now, + ) + + +def _launch_request( + *, + suite_version: str | None, + first_origin: str | None = "stored", + first_version: str | None = "first-v1", + second_origin: str | None = "stored", + second_version: str | None = "second-v1", +) -> LaunchScenarioSuiteRequest: + """Build the reviewed suite and preset snapshot sent by the UI.""" + return LaunchScenarioSuiteRequest( + expected_version=suite_version, + entries=[ + ScenarioSuiteLaunchEntrySnapshot( + entry_id=ENTRY_ONE, + preset_name="first", + origin=first_origin, + preset_version=first_version, + ), + ScenarioSuiteLaunchEntrySnapshot( + entry_id=ENTRY_TWO, + preset_name="second", + origin=second_origin, + preset_version=second_version, + ), + ], + target_name="objective-target", + adversarial_target_name="adversarial-target", + max_concurrency=7, + max_retries=3, + operation_id=OPERATION_ID, + labels={"team": "purple"}, + ) + + +async def _wait_for_submissions(service: ScenarioSuiteExecutionService) -> None: + """Wait for immediate mock submissions without sleeping wall-clock time.""" + for _ in range(100): + if not service.has_active_work(): + return + await asyncio.sleep(0) + raise AssertionError("Suite submission coordinator did not finish") + + +@pytest.fixture +async def execution_context( + tmp_path: Path, +) -> AsyncGenerator[ + tuple[ + ScenarioSuiteExecutionService, + ScenarioSuiteRegistry, + MagicMock, + MagicMock, + list[tuple[RunScenarioRequest, dict[str, Any]]], + dict[str, ScenarioRunSummary], + ], + None, +]: + """Create isolated registries and a controllable run service.""" + suite_registry = ScenarioSuiteRegistry(storage=ScenarioSuiteStorage(source=str(tmp_path))) + preset_registry = MagicMock(spec=ScenarioPresetRegistry) + presets = { + "first": _preset_entry( + name="first", + scenario_name="scenario.first", + version="first-v1", + ), + "second": _preset_entry( + name="second", + scenario_name="scenario.second", + version="second-v1", + ), + } + preset_registry.get_preset.side_effect = lambda *, name: presets.get(name) + preset_registry.check_preset.return_value = [] + + calls: list[tuple[RunScenarioRequest, dict[str, Any]]] = [] + runs: dict[str, ScenarioRunSummary] = {} + + async def start_run_async( + *, + request: RunScenarioRequest, + initial_metadata: dict[str, Any], + ) -> ScenarioRunSummary: + calls.append((request, initial_metadata)) + run_id = f"run-{len(calls)}" + summary = _run_summary(run_id=run_id, scenario_name=request.scenario_name) + runs[run_id] = summary + return summary + + async def get_run_async(*, scenario_result_id: str) -> ScenarioRunSummary | None: + return runs.get(scenario_result_id) + + async def cancel_run_async(*, scenario_result_id: str) -> ScenarioRunSummary | None: + current = runs.get(scenario_result_id) + if current is None: + return None + cancelled = current.model_copy(update={"status": ScenarioRunState.CANCELLED}) + runs[scenario_result_id] = cancelled + return cancelled + + run_service = MagicMock(spec=ScenarioRunService) + run_service.start_run_async = AsyncMock(side_effect=start_run_async) + run_service.get_run_async = AsyncMock(side_effect=get_run_async) + run_service.cancel_run_async = AsyncMock(side_effect=cancel_run_async) + + service = ScenarioSuiteExecutionService( + suite_registry=suite_registry, + preset_registry=preset_registry, + run_service=run_service, + ) + try: + yield service, suite_registry, preset_registry, run_service, calls, runs + finally: + await service.shutdown_async() + + +async def test_launch_queues_entries_in_order_and_stamps_versioned_metadata( + execution_context: tuple[ + ScenarioSuiteExecutionService, + ScenarioSuiteRegistry, + MagicMock, + MagicMock, + list[tuple[RunScenarioRequest, dict[str, Any]]], + dict[str, ScenarioRunSummary], + ], +) -> None: + """Test server-owned identity, launch propagation, ordering, and child metadata.""" + service, suite_registry, _, _, calls, _ = execution_context + stored = suite_registry.save_suite(suite=_suite(), expected_version=None) + + accepted = await service.launch_async( + name="weekly", + request=_launch_request(suite_version=stored.version), + ) + await _wait_for_submissions(service) + summary = await service.get_execution_async(execution_id=accepted.suite_execution_id) + + assert summary is not None + assert summary.status == ScenarioSuiteExecutionStatus.IN_PROGRESS + assert [entry.preset_name for entry in summary.entries] == ["first", "second"] + assert [entry.entry_position for entry in summary.entries] == [1, 2] + assert [request.scenario_name for request, _ in calls] == [ + "scenario.first", + "scenario.second", + ] + for request, _ in calls: + assert request.target_name == "objective-target" + assert request.adversarial_target_name == "adversarial-target" + assert request.max_concurrency == 7 + assert request.max_retries == 3 + assert request.operation_id == OPERATION_ID + assert request.labels == {"team": "purple"} + + first_metadata = calls[0][1][SCENARIO_SUITE_EXECUTION_METADATA_KEY] + second_metadata = calls[1][1][SCENARIO_SUITE_EXECUTION_METADATA_KEY] + assert first_metadata == { + "version": 1, + "suite_execution_id": str(accepted.suite_execution_id), + "suite_name": "weekly", + "suite_version": stored.version, + "suite_entry_id": str(ENTRY_ONE), + "entry_position": 1, + "preset_name": "first", + "preset_version": "first-v1", + } + assert second_metadata["suite_execution_id"] == str(accepted.suite_execution_id) + assert second_metadata["suite_entry_id"] == str(ENTRY_TWO) + assert second_metadata["entry_position"] == 2 + + +async def test_registered_suite_and_presets_launch_with_nullable_versions( + execution_context: tuple[ + ScenarioSuiteExecutionService, + ScenarioSuiteRegistry, + MagicMock, + MagicMock, + list[tuple[RunScenarioRequest, dict[str, Any]]], + dict[str, ScenarioRunSummary], + ], +) -> None: + """Test code-defined definitions use the same launch path with null versions.""" + service, suite_registry, preset_registry, _, calls, _ = execution_context + suite_registry.register_suite(suite=_suite()) + preset_registry.get_preset.side_effect = lambda *, name: _preset_entry( + name=name, + scenario_name=f"scenario.{name}", + origin="registered", + version=None, + ) + + accepted = await service.launch_async( + name="weekly", + request=_launch_request( + suite_version=None, + first_origin="registered", + first_version=None, + second_origin="registered", + second_version=None, + ), + ) + await _wait_for_submissions(service) + + summary = await service.get_execution_async(execution_id=accepted.suite_execution_id) + assert summary is not None + assert summary.suite_origin == "registered" + assert summary.suite_version is None + assert [entry.preset_version for entry in summary.entries] == [None, None] + assert calls[0][1][SCENARIO_SUITE_EXECUTION_METADATA_KEY]["suite_version"] is None + assert calls[0][1][SCENARIO_SUITE_EXECUTION_METADATA_KEY]["preset_version"] is None + + +async def test_stored_suite_requires_the_reviewed_version( + execution_context: tuple[ + ScenarioSuiteExecutionService, + ScenarioSuiteRegistry, + MagicMock, + MagicMock, + list[tuple[RunScenarioRequest, dict[str, Any]]], + dict[str, ScenarioRunSummary], + ], +) -> None: + """Test a stored suite cannot silently launch without optimistic concurrency.""" + service, suite_registry, _, run_service, _, _ = execution_context + suite_registry.save_suite(suite=_suite(), expected_version=None) + + with pytest.raises(ScenarioSuiteLaunchConflictError, match="no reviewed version"): + await service.launch_async(name="weekly", request=_launch_request(suite_version=None)) + + run_service.start_run_async.assert_not_awaited() + + +async def test_changed_preset_snapshot_rejects_the_whole_launch_before_submission( + execution_context: tuple[ + ScenarioSuiteExecutionService, + ScenarioSuiteRegistry, + MagicMock, + MagicMock, + list[tuple[RunScenarioRequest, dict[str, Any]]], + dict[str, ScenarioRunSummary], + ], +) -> None: + """Test referenced preset drift is a hard preflight conflict.""" + service, suite_registry, _, run_service, _, _ = execution_context + stored = suite_registry.save_suite(suite=_suite(), expected_version=None) + + with pytest.raises(ScenarioSuiteLaunchConflictError, match="first"): + await service.launch_async( + name="weekly", + request=_launch_request(suite_version=stored.version, first_version="stale"), + ) + + run_service.start_run_async.assert_not_awaited() + + +async def test_advisory_preset_issue_becomes_a_hard_launch_gate( + execution_context: tuple[ + ScenarioSuiteExecutionService, + ScenarioSuiteRegistry, + MagicMock, + MagicMock, + list[tuple[RunScenarioRequest, dict[str, Any]]], + dict[str, ScenarioRunSummary], + ], +) -> None: + """Test suite editing remains permissive while execution rejects invalid entries.""" + service, suite_registry, preset_registry, run_service, _, _ = execution_context + stored = suite_registry.save_suite(suite=_suite(), expected_version=None) + preset_registry.check_preset.side_effect = lambda *, preset: ( + [PresetIssue(field="scenario_name", message="Scenario is unavailable.")] if preset.name == "second" else [] + ) + + with pytest.raises(ScenarioSuiteNotRunnableError, match="Entry 2"): + await service.launch_async( + name="weekly", + request=_launch_request(suite_version=stored.version), + ) + + run_service.start_run_async.assert_not_awaited() + + +async def test_submission_failure_does_not_prevent_later_entries_from_queueing( + execution_context: tuple[ + ScenarioSuiteExecutionService, + ScenarioSuiteRegistry, + MagicMock, + MagicMock, + list[tuple[RunScenarioRequest, dict[str, Any]]], + dict[str, ScenarioRunSummary], + ], +) -> None: + """Test accepted entries retain independent ordinary-run failure semantics.""" + service, suite_registry, _, run_service, calls, runs = execution_context + stored = suite_registry.save_suite(suite=_suite(), expected_version=None) + + async def start_run_async( + *, + request: RunScenarioRequest, + initial_metadata: dict[str, Any], + ) -> ScenarioRunSummary: + calls.append((request, initial_metadata)) + if len(calls) == 1: + raise ValueError("first entry could not initialize") + summary = _run_summary( + run_id="run-2", + scenario_name=request.scenario_name, + state=ScenarioRunState.COMPLETED, + ) + runs[summary.scenario_result_id] = summary + return summary + + run_service.start_run_async.side_effect = start_run_async + + accepted = await service.launch_async( + name="weekly", + request=_launch_request(suite_version=stored.version), + ) + await _wait_for_submissions(service) + summary = await service.get_execution_async(execution_id=accepted.suite_execution_id) + + assert summary is not None + assert summary.status == ScenarioSuiteExecutionStatus.FAILED + assert [entry.state for entry in summary.entries] == [ + ScenarioSuiteExecutionEntryState.SUBMISSION_FAILED, + ScenarioSuiteExecutionEntryState.COMPLETED, + ] + assert summary.entries[0].error == "first entry could not initialize" + assert summary.entries[1].run is not None + assert len(calls) == 2 + + +async def test_cancel_stops_pending_submission_and_cancels_created_children( + execution_context: tuple[ + ScenarioSuiteExecutionService, + ScenarioSuiteRegistry, + MagicMock, + MagicMock, + list[tuple[RunScenarioRequest, dict[str, Any]]], + dict[str, ScenarioRunSummary], + ], +) -> None: + """Test best-effort whole-execution cancellation without a persisted parent.""" + service, suite_registry, _, run_service, calls, runs = execution_context + stored = suite_registry.save_suite(suite=_suite(), expected_version=None) + hold_second = asyncio.Event() + + async def start_run_async( + *, + request: RunScenarioRequest, + initial_metadata: dict[str, Any], + ) -> ScenarioRunSummary: + calls.append((request, initial_metadata)) + if len(calls) == 2: + await hold_second.wait() + summary = _run_summary( + run_id=f"run-{len(calls)}", + scenario_name=request.scenario_name, + ) + runs[summary.scenario_result_id] = summary + return summary + + run_service.start_run_async.side_effect = start_run_async + + accepted = await service.launch_async( + name="weekly", + request=_launch_request(suite_version=stored.version), + ) + for _ in range(100): + current = await service.get_execution_async(execution_id=accepted.suite_execution_id) + if current is not None and current.entries[0].run is not None: + break + await asyncio.sleep(0) + else: + raise AssertionError("First suite child was not submitted") + + await service.cancel_execution_async(execution_id=accepted.suite_execution_id) + await _wait_for_submissions(service) + summary = await service.get_execution_async(execution_id=accepted.suite_execution_id) + + assert summary is not None + assert summary.status == ScenarioSuiteExecutionStatus.CANCELLED + assert [entry.state for entry in summary.entries] == [ + ScenarioSuiteExecutionEntryState.CANCELLED, + ScenarioSuiteExecutionEntryState.CANCELLED, + ] + run_service.cancel_run_async.assert_awaited_once_with(scenario_result_id="run-1") diff --git a/tests/unit/backend/test_scenario_suite_routes.py b/tests/unit/backend/test_scenario_suite_routes.py index 0b911fcba5..c10cc75a35 100644 --- a/tests/unit/backend/test_scenario_suite_routes.py +++ b/tests/unit/backend/test_scenario_suite_routes.py @@ -4,6 +4,7 @@ """Tests for scenario suite definition routes.""" from collections.abc import Iterator +from datetime import UTC, datetime from pathlib import Path from typing import Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -17,9 +18,19 @@ from pyrit.backend.middleware.auth import current_user_name, require_admin from pyrit.backend.models.scenario_suites import ( ScenarioSuiteEntryStatus, + ScenarioSuiteExecutionEntryState, + ScenarioSuiteExecutionEntrySummary, + ScenarioSuiteExecutionStatus, + ScenarioSuiteExecutionSummary, ScenarioSuiteListResponse, ScenarioSuiteResponse, ) +from pyrit.backend.services.scenario_suite_execution_service import ( + ScenarioSuiteExecutionNotFoundError, + ScenarioSuiteExecutionService, + ScenarioSuiteLaunchConflictError, + ScenarioSuiteNotRunnableError, +) from pyrit.backend.services.scenario_suite_service import ScenarioSuiteService from pyrit.models.catalog import PresetIssue, ScenarioSuite, ScenarioSuiteEntry, ScenarioSuiteRegistryEntry from pyrit.registry import ( @@ -32,6 +43,7 @@ SUITE_NAME = "weekly" ENTRY_ID = UUID("11111111-1111-4111-8111-111111111111") +EXECUTION_ID = UUID("eeeeeeee-eeee-4eee-8eee-eeeeeeeeeeee") def _suite(*, name: str = SUITE_NAME, author: str | None = None) -> ScenarioSuite: @@ -68,6 +80,51 @@ def _response( ) +def _execution_summary( + *, + status_value: ScenarioSuiteExecutionStatus = ScenarioSuiteExecutionStatus.QUEUING, +) -> ScenarioSuiteExecutionSummary: + """Build one transient suite execution response.""" + now = datetime(2025, 1, 1, tzinfo=UTC) + return ScenarioSuiteExecutionSummary( + suite_execution_id=EXECUTION_ID, + suite_name=SUITE_NAME, + suite_origin="stored", + suite_version="suite-v1", + status=status_value, + created_at=now, + updated_at=now, + entries=[ + ScenarioSuiteExecutionEntrySummary( + entry_id=ENTRY_ID, + entry_position=1, + preset_name="nightly", + preset_origin="stored", + preset_version="preset-v1", + state=ScenarioSuiteExecutionEntryState.PENDING, + ) + ], + ) + + +def _launch_body() -> dict[str, object]: + """Build the suite launch payload sent by the browser.""" + return { + "expected_version": "suite-v1", + "entries": [ + { + "entry_id": str(ENTRY_ID), + "preset_name": "nightly", + "origin": "stored", + "preset_version": "preset-v1", + } + ], + "target_name": "objective-target", + "max_concurrency": 5, + "max_retries": 2, + } + + @pytest.fixture def client(compatibility_headers: dict[str, str]) -> Iterator[TestClient]: """Create a client with administrator authorization satisfied.""" @@ -98,6 +155,17 @@ def service() -> Iterator[MagicMock]: yield mock_service +@pytest.fixture +def execution_service() -> Iterator[MagicMock]: + """Patch the transient suite execution coordinator resolved by the routes.""" + mock_service = MagicMock(spec=ScenarioSuiteExecutionService) + with patch( + "pyrit.backend.routes.scenario_suites.get_scenario_suite_execution_service", + return_value=mock_service, + ): + yield mock_service + + @pytest.fixture def unavailable_preset_backend( tmp_path: Path, @@ -141,6 +209,94 @@ def test_list_is_open_and_returns_entry_issues(anonymous_client: TestClient, ser assert response.json()["items"][0]["entry_statuses"][0]["issues"][0]["message"] == "Preset is missing." +def test_launch_is_open_and_returns_server_execution_id( + anonymous_client: TestClient, + execution_service: MagicMock, +) -> None: + """Test one launch action creates a server-owned transient execution.""" + execution_service.launch_async = AsyncMock(return_value=_execution_summary()) + + response = anonymous_client.post(f"/api/scenario-suites/{SUITE_NAME}/launch", json=_launch_body()) + + assert response.status_code == 202 + assert response.json()["suite_execution_id"] == str(EXECUTION_ID) + launch = execution_service.launch_async.call_args.kwargs["request"] + assert launch.expected_version == "suite-v1" + assert launch.entries[0].preset_version == "preset-v1" + assert launch.max_concurrency == 5 + + +def test_launch_is_rejected_during_runtime_reinitialization( + anonymous_client: TestClient, + execution_service: MagicMock, +) -> None: + """Test runtime transitions reject new suite work before preflight.""" + with patch.object(app.state, "runtime_lifecycle", MagicMock(state="initializing"), create=True): + response = anonymous_client.post( + f"/api/scenario-suites/{SUITE_NAME}/launch", + json=_launch_body(), + ) + + assert response.status_code == 503 + execution_service.launch_async.assert_not_called() + + +@pytest.mark.parametrize( + ("error", "expected_status"), + [ + (ScenarioSuiteExecutionNotFoundError("missing"), 404), + (ScenarioSuiteLaunchConflictError("changed"), 409), + (ScenarioSuiteNotRunnableError("not runnable"), 400), + (AzureError("storage unavailable"), 503), + (RuntimeError("scheduling is stopping"), 503), + ], +) +def test_launch_maps_preflight_failures( + anonymous_client: TestClient, + execution_service: MagicMock, + error: Exception, + expected_status: int, +) -> None: + """Test actionable launch failures without exposing storage details.""" + execution_service.launch_async = AsyncMock(side_effect=error) + + response = anonymous_client.post(f"/api/scenario-suites/{SUITE_NAME}/launch", json=_launch_body()) + + assert response.status_code == expected_status + if isinstance(error, AzureError): + assert "storage unavailable" not in response.text + + +def test_get_and_cancel_transient_execution( + anonymous_client: TestClient, + execution_service: MagicMock, +) -> None: + """Test grouped progress reads and best-effort cancellation.""" + execution_service.get_execution_async = AsyncMock(return_value=_execution_summary()) + execution_service.cancel_execution_async = AsyncMock( + return_value=_execution_summary(status_value=ScenarioSuiteExecutionStatus.CANCELLED) + ) + + get_response = anonymous_client.get(f"/api/scenario-suites/executions/{EXECUTION_ID}") + cancel_response = anonymous_client.post(f"/api/scenario-suites/executions/{EXECUTION_ID}/cancel") + + assert get_response.status_code == 200 + assert cancel_response.status_code == 200 + assert cancel_response.json()["status"] == "CANCELLED" + + +def test_missing_transient_execution_returns_not_found( + anonymous_client: TestClient, + execution_service: MagicMock, +) -> None: + """Test process-local progress expiry is explicit.""" + execution_service.get_execution_async = AsyncMock(return_value=None) + execution_service.cancel_execution_async = AsyncMock(return_value=None) + + assert anonymous_client.get(f"/api/scenario-suites/executions/{EXECUTION_ID}").status_code == 404 + assert anonymous_client.post(f"/api/scenario-suites/executions/{EXECUTION_ID}/cancel").status_code == 404 + + def test_get_returns_suite_and_version(client: TestClient, service: MagicMock) -> None: """Test a single open read.""" service.get_suite_async = AsyncMock(return_value=_response()) diff --git a/tests/unit/models/test_scenario_suite_execution.py b/tests/unit/models/test_scenario_suite_execution.py new file mode 100644 index 0000000000..06ac0bee4c --- /dev/null +++ b/tests/unit/models/test_scenario_suite_execution.py @@ -0,0 +1,34 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Tests for versioned scenario suite execution metadata.""" + +from uuid import UUID + +from pyrit.models import ( + SCENARIO_SUITE_EXECUTION_METADATA_KEY, + ScenarioSuiteExecutionMetadata, +) + + +def test_suite_execution_metadata_round_trips_through_scenario_metadata() -> None: + """Test the persisted envelope keeps server and definition identities intact.""" + metadata = ScenarioSuiteExecutionMetadata( + suite_execution_id=UUID("eeeeeeee-eeee-4eee-8eee-eeeeeeeeeeee"), + suite_name="weekly", + suite_version="suite-v1", + suite_entry_id=UUID("11111111-1111-4111-8111-111111111111"), + entry_position=2, + preset_name="nightly", + preset_version="preset-v1", + ) + + fields = metadata.to_metadata_fields() + + assert fields[SCENARIO_SUITE_EXECUTION_METADATA_KEY]["version"] == 1 + assert ScenarioSuiteExecutionMetadata.from_metadata(fields) == metadata + + +def test_direct_run_has_no_suite_execution_metadata() -> None: + """Test absence remains distinct from an invalid or empty suite envelope.""" + assert ScenarioSuiteExecutionMetadata.from_metadata({"run_plan": {}}) is None From d61be4a6ffa87db01eb39359dddf63b56759a3bf Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Fri, 23 Oct 2026 12:04:09 -0400 Subject: [PATCH 4/7] FIX: Sanitize ScenarioSuite execution errors Keep full exception details in server logs while exposing only stable user-safe failure messages through suite execution responses and the UI. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 2403a4a9-6af6-46b2-b371-46de73a018fd --- .../ScenarioSuiteExecutionPage.test.tsx | 27 ++++++++-- .../ScenarioSuiteExecutionPage.tsx | 15 +++++- frontend/src/types/index.ts | 2 + pyrit/backend/models/scenario_suites.py | 10 +++- .../scenario_suite_execution_service.py | 13 ++++- .../test_scenario_suite_execution_service.py | 49 ++++++++++++++++++- 6 files changed, 105 insertions(+), 11 deletions(-) diff --git a/frontend/src/components/ScenarioSuites/ScenarioSuiteExecutionPage.test.tsx b/frontend/src/components/ScenarioSuites/ScenarioSuiteExecutionPage.test.tsx index fea7c27334..3b55fe022a 100644 --- a/frontend/src/components/ScenarioSuites/ScenarioSuiteExecutionPage.test.tsx +++ b/frontend/src/components/ScenarioSuites/ScenarioSuiteExecutionPage.test.tsx @@ -172,7 +172,12 @@ describe('ScenarioSuiteExecutionPage', () => { expect(within(rows[1]).getAllByText('Queued')).toHaveLength(2) }) - it('surfaces entry submission failures without inventing a child run', () => { + it('shows safe entry failures without rendering backend exception details', () => { + const rawError = ( + "(sqlite3.OperationalError) database is locked " + + "[SQL: INSERT INTO ScenarioResultEntries ...] " + + "[parameters: ('secret-token',)]" + ) setHookResult({ loadStatus: 'ready', execution: makeExecution({ @@ -183,7 +188,14 @@ describe('ScenarioSuiteExecutionPage', () => { position: 1, presetName: 'first_preset', state: 'SUBMISSION_FAILED', - error: 'Scenario could not initialize.', + error: rawError, + }), + makeEntry({ + entryId: SECOND_ENTRY_ID, + position: 2, + presetName: 'second_preset', + state: 'FAILED', + error: rawError, }), ], }), @@ -193,8 +205,15 @@ describe('ScenarioSuiteExecutionPage', () => { renderPage() - expect(screen.getByText('Scenario could not initialize.')).toBeInTheDocument() - expect(screen.getByText('Not created')).toBeInTheDocument() + expect(screen.getByText( + 'The scenario run could not be created. Check server logs before retrying.', + )).toBeInTheDocument() + expect(screen.getByText( + 'The scenario run status could not be loaded. Refresh the page or check server logs.', + )).toBeInTheDocument() + expect(screen.queryByText(/database is locked/i)).not.toBeInTheDocument() + expect(screen.queryByText(/secret-token/i)).not.toBeInTheDocument() + expect(screen.getAllByText('Not created')).toHaveLength(2) expect(screen.getAllByText('Submission failed')).toHaveLength(2) }) diff --git a/frontend/src/components/ScenarioSuites/ScenarioSuiteExecutionPage.tsx b/frontend/src/components/ScenarioSuites/ScenarioSuiteExecutionPage.tsx index 2f8e0a5013..57504789f2 100644 --- a/frontend/src/components/ScenarioSuites/ScenarioSuiteExecutionPage.tsx +++ b/frontend/src/components/ScenarioSuites/ScenarioSuiteExecutionPage.tsx @@ -74,6 +74,13 @@ const ENTRY_BADGE_COLORS: Record< CANCELLED: 'warning', } +const ENTRY_SUBMISSION_FAILURE_MESSAGE = ( + 'The scenario run could not be created. Check server logs before retrying.' +) +const ENTRY_STATUS_FAILURE_MESSAGE = ( + 'The scenario run status could not be loaded. Refresh the page or check server logs.' +) + export default function ScenarioSuiteExecutionPage() { const { executionId } = useParams<{ executionId: string }>() const decodedExecutionId = routerPathParamValue(executionId) @@ -300,7 +307,7 @@ function ScenarioSuiteExecutionPageContent({ {entry.error && ( - {entry.error} + {entryErrorMessage(entry)} )}
@@ -408,6 +415,12 @@ function entryProgress(entry: ScenarioSuiteExecutionEntrySummary): string { return `${entry.run.completed_attacks} of ${entry.run.total_attacks} attacks` } +function entryErrorMessage(entry: ScenarioSuiteExecutionEntrySummary): string { + return entry.state === 'SUBMISSION_FAILED' + ? ENTRY_SUBMISSION_FAILURE_MESSAGE + : ENTRY_STATUS_FAILURE_MESSAGE +} + function statusIcon( status: ScenarioSuiteExecutionStatus | ScenarioSuiteExecutionEntryState, ): ReactElement | undefined { diff --git a/frontend/src/types/index.ts b/frontend/src/types/index.ts index f6431d4524..a9d5a52ef9 100644 --- a/frontend/src/types/index.ts +++ b/frontend/src/types/index.ts @@ -1498,7 +1498,9 @@ export interface ScenarioSuiteExecutionEntrySummary { preset_version: string | null state: ScenarioSuiteExecutionEntryState run: ScenarioRunSummary | null + /** User-safe failure text. The UI still maps it to state-owned copy before display. */ error: string | null + /** Internal exception class for diagnostics; never display it to operators. */ error_type: string | null } diff --git a/pyrit/backend/models/scenario_suites.py b/pyrit/backend/models/scenario_suites.py index f2de4d219b..91cd9c6c81 100644 --- a/pyrit/backend/models/scenario_suites.py +++ b/pyrit/backend/models/scenario_suites.py @@ -187,8 +187,14 @@ class ScenarioSuiteExecutionEntrySummary(BaseModel): ) state: ScenarioSuiteExecutionEntryState = Field(..., description="Submission or child-run state") run: ScenarioRunSummary | None = Field(None, description="Current child run summary after successful submission") - error: str | None = Field(None, description="Submission failure when no child run could be created") - error_type: str | None = Field(None, description="Submission exception class name") + error: str | None = Field( + None, + description="User-safe submission or status failure message when no child run can be shown", + ) + error_type: str | None = Field( + None, + description="Internal exception class name for diagnostics; clients should not display it", + ) class ScenarioSuiteExecutionSummary(BaseModel): diff --git a/pyrit/backend/services/scenario_suite_execution_service.py b/pyrit/backend/services/scenario_suite_execution_service.py index 157e9e4e00..735f74bc53 100644 --- a/pyrit/backend/services/scenario_suite_execution_service.py +++ b/pyrit/backend/services/scenario_suite_execution_service.py @@ -127,6 +127,9 @@ class ScenarioSuiteExecutionService: remains an ordinary persisted ``ScenarioResult`` managed by ``ScenarioRunService``. """ + SUBMISSION_FAILURE_MESSAGE = "The scenario run could not be created. Check server logs before retrying." + CHILD_STATUS_FAILURE_MESSAGE = "The scenario run status could not be loaded. Refresh the page or check server logs." + def __init__( self, *, @@ -382,7 +385,7 @@ async def _submit_entries_async(self, *, state: _ExecutionState) -> None: ) async with state.lock: entry.state = ScenarioSuiteExecutionEntryState.SUBMISSION_FAILED - entry.error = str(exc) + entry.error = self.SUBMISSION_FAILURE_MESSAGE entry.error_type = type(exc).__name__ state.updated_at = datetime.now(UTC) continue @@ -473,8 +476,14 @@ async def _build_summary_from_snapshot_async( error_type = entry.error_type entry_state = entry.state if isinstance(run_or_error, BaseException): + logger.warning( + "Scenario suite execution %s entry %s child run status could not be loaded.", + snapshot.execution_id, + entry.resolved.entry_id, + exc_info=(type(run_or_error), run_or_error, run_or_error.__traceback__), + ) entry_state = ScenarioSuiteExecutionEntryState.FAILED - error = str(run_or_error) + error = self.CHILD_STATUS_FAILURE_MESSAGE error_type = type(run_or_error).__name__ elif run_or_error is not None: run = run_or_error diff --git a/tests/unit/backend/test_scenario_suite_execution_service.py b/tests/unit/backend/test_scenario_suite_execution_service.py index a2d195c64d..18715f9b4d 100644 --- a/tests/unit/backend/test_scenario_suite_execution_service.py +++ b/tests/unit/backend/test_scenario_suite_execution_service.py @@ -386,7 +386,11 @@ async def start_run_async( ) -> ScenarioRunSummary: calls.append((request, initial_metadata)) if len(calls) == 1: - raise ValueError("first entry could not initialize") + raise ValueError( + "(sqlite3.OperationalError) database is locked " + "[SQL: INSERT INTO ScenarioResultEntries ...] " + "[parameters: ('secret-token',)]" + ) summary = _run_summary( run_id="run-2", scenario_name=request.scenario_name, @@ -410,11 +414,52 @@ async def start_run_async( ScenarioSuiteExecutionEntryState.SUBMISSION_FAILED, ScenarioSuiteExecutionEntryState.COMPLETED, ] - assert summary.entries[0].error == "first entry could not initialize" + assert summary.entries[0].error == service.SUBMISSION_FAILURE_MESSAGE + assert summary.entries[0].error_type == "ValueError" + assert "database is locked" not in summary.model_dump_json() + assert "secret-token" not in summary.model_dump_json() assert summary.entries[1].run is not None assert len(calls) == 2 +async def test_child_status_failure_does_not_expose_exception_details( + execution_context: tuple[ + ScenarioSuiteExecutionService, + ScenarioSuiteRegistry, + MagicMock, + MagicMock, + list[tuple[RunScenarioRequest, dict[str, Any]]], + dict[str, ScenarioRunSummary], + ], +) -> None: + """Test child status lookup failures retain diagnostics without exposing exception text.""" + service, suite_registry, _, run_service, _, runs = execution_context + stored = suite_registry.save_suite(suite=_suite(), expected_version=None) + + accepted = await service.launch_async( + name="weekly", + request=_launch_request(suite_version=stored.version), + ) + await _wait_for_submissions(service) + + async def get_run_async(*, scenario_result_id: str) -> ScenarioRunSummary | None: + if scenario_result_id == "run-1": + raise RuntimeError("SELECT secret_column FROM private_table") + return runs.get(scenario_result_id) + + runs["run-2"] = runs["run-2"].model_copy(update={"status": ScenarioRunState.COMPLETED}) + run_service.get_run_async.side_effect = get_run_async + summary = await service.get_execution_async(execution_id=accepted.suite_execution_id) + + assert summary is not None + assert summary.status == ScenarioSuiteExecutionStatus.FAILED + assert summary.entries[0].state == ScenarioSuiteExecutionEntryState.FAILED + assert summary.entries[0].error == service.CHILD_STATUS_FAILURE_MESSAGE + assert summary.entries[0].error_type == "RuntimeError" + assert "secret_column" not in summary.model_dump_json() + assert summary.entries[1].run is not None + + async def test_cancel_stops_pending_submission_and_cancels_created_children( execution_context: tuple[ ScenarioSuiteExecutionService, From fda15389c39d99348c3b508f816efa66203d9ae8 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Fri, 23 Oct 2026 12:15:44 -0400 Subject: [PATCH 5/7] TEST: Forward ScenarioRun initial metadata in mocks Keep patched ScenarioRunService helpers compatible with the internal initial_metadata keyword added for suite child attribution. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 2403a4a9-6af6-46b2-b371-46de73a018fd --- tests/unit/backend/test_scenario_resume.py | 5 +++-- tests/unit/backend/test_technique_service.py | 11 ++++++++--- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/tests/unit/backend/test_scenario_resume.py b/tests/unit/backend/test_scenario_resume.py index d234a450a5..55c52ef4c9 100644 --- a/tests/unit/backend/test_scenario_resume.py +++ b/tests/unit/backend/test_scenario_resume.py @@ -7,7 +7,7 @@ import pathlib import threading import uuid -from collections.abc import AsyncIterator, Awaitable, Sequence +from collections.abc import AsyncIterator, Awaitable, Mapping, Sequence from typing import Any, ClassVar from unittest.mock import AsyncMock, MagicMock, patch @@ -868,7 +868,8 @@ async def test_overlapping_resume_requests_only_initialize_once_async( entered = asyncio.Event() release = asyncio.Event() - async def fail_launch_async(*, request: RunScenarioRequest) -> None: + async def fail_launch_async(*, request: RunScenarioRequest, initial_metadata: Mapping[str, Any]) -> None: + assert initial_metadata == {} entered.set() await release.wait() raise ValueError("Synthetic initialization failure") diff --git a/tests/unit/backend/test_technique_service.py b/tests/unit/backend/test_technique_service.py index ce3cf9c570..ce892dfc95 100644 --- a/tests/unit/backend/test_technique_service.py +++ b/tests/unit/backend/test_technique_service.py @@ -6,7 +6,7 @@ import ast import asyncio import threading -from collections.abc import Iterator +from collections.abc import Iterator, Mapping from enum import Enum from typing import Any from unittest.mock import AsyncMock, patch @@ -457,9 +457,14 @@ async def test_launch_survives_registration_between_resolution_and_construction_ service = ScenarioRunService() initialize = service._initialize_scenario_async - async def register_then_initialize_async(*, request: RunScenarioRequest, init_kwargs: dict[str, Any]) -> Scenario: + async def register_then_initialize_async( + *, + request: RunScenarioRequest, + init_kwargs: dict[str, Any], + initial_metadata: Mapping[str, Any] | None = None, + ) -> Scenario: await TechniqueService().create_async(CreateTechniqueRequest(name="during_launch", type="PromptSendingAttack")) - return await initialize(request=request, init_kwargs=init_kwargs) + return await initialize(request=request, init_kwargs=init_kwargs, initial_metadata=initial_metadata) try: with ( From b3c3da1574bed07b6c16b62b2e1417810d977011 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Fri, 23 Oct 2026 12:39:48 -0400 Subject: [PATCH 6/7] FIX: Harden ScenarioSuite execution lifecycle Preserve admitted child runs across suite cancellation, keep transient status-read failures nonterminal, and reject capacity when only active executions remain. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 2403a4a9-6af6-46b2-b371-46de73a018fd --- pyrit/backend/routes/scenario_suites.py | 4 + .../scenario_suite_execution_service.py | 264 +++++++++++------- .../test_scenario_suite_execution_service.py | 165 +++++++++-- .../backend/test_scenario_suite_routes.py | 4 + 4 files changed, 312 insertions(+), 125 deletions(-) diff --git a/pyrit/backend/routes/scenario_suites.py b/pyrit/backend/routes/scenario_suites.py index 07eeb4551f..fc95076360 100644 --- a/pyrit/backend/routes/scenario_suites.py +++ b/pyrit/backend/routes/scenario_suites.py @@ -19,6 +19,7 @@ UpdateScenarioSuiteRequest, ) from pyrit.backend.services.scenario_suite_execution_service import ( + ScenarioSuiteExecutionCapacityError, ScenarioSuiteExecutionNotFoundError, ScenarioSuiteLaunchConflictError, ScenarioSuiteNotRunnableError, @@ -144,6 +145,7 @@ async def cancel_scenario_suite_execution( # pyrit-async-suffix-exempt 400: {"model": ProblemDetail, "description": "The suite has unresolved or invalid entries"}, 404: {"model": ProblemDetail, "description": "Scenario suite not found"}, 409: {"model": ProblemDetail, "description": "The suite or a referenced preset changed"}, + 429: {"model": ProblemDetail, "description": "Active suite execution capacity is full"}, 503: {"model": ProblemDetail, "description": "Runtime or definition storage unavailable"}, }, ) @@ -178,6 +180,8 @@ async def launch_scenario_suite( # pyrit-async-suffix-exempt ) from exc except ScenarioSuiteNotRunnableError as exc: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None + except ScenarioSuiteExecutionCapacityError as exc: + raise HTTPException(status_code=status.HTTP_429_TOO_MANY_REQUESTS, detail=str(exc)) from None except AzureError as exc: raise _storage_unavailable() from exc except RuntimeError as exc: diff --git a/pyrit/backend/services/scenario_suite_execution_service.py b/pyrit/backend/services/scenario_suite_execution_service.py index 735f74bc53..db4c98e5ae 100644 --- a/pyrit/backend/services/scenario_suite_execution_service.py +++ b/pyrit/backend/services/scenario_suite_execution_service.py @@ -32,12 +32,26 @@ logger = logging.getLogger(__name__) _MAX_RETAINED_EXECUTIONS = 256 +_CAPACITY_MESSAGE = ( + "Too many scenario suites are still active. Wait for an active suite to finish before launching another." +) +_TERMINAL_EXECUTION_STATUSES = frozenset( + { + ScenarioSuiteExecutionStatus.COMPLETED, + ScenarioSuiteExecutionStatus.FAILED, + ScenarioSuiteExecutionStatus.CANCELLED, + } +) class ScenarioSuiteExecutionNotFoundError(LookupError): """A requested suite definition or transient execution does not exist.""" +class ScenarioSuiteExecutionCapacityError(RuntimeError): + """Active suite executions fill the transient coordinator capacity.""" + + class ScenarioSuiteLaunchConflictError(ValueError): """A suite or referenced preset changed after the operator reviewed it.""" @@ -143,6 +157,7 @@ def __init__( self._run_service = run_service self._executions: OrderedDict[UUID, _ExecutionState] = OrderedDict() self._tasks: set[asyncio.Task[None]] = set() + self._admission_lock = asyncio.Lock() self._stopping = False def has_active_work(self) -> bool: @@ -177,20 +192,24 @@ async def launch_async( if self._stopping: raise RuntimeError("Scenario suite launch scheduling is stopping.") - now = datetime.now(UTC) - state = _ExecutionState( - execution_id=uuid4(), - suite=suite, - entries=[_ExecutionEntryState(resolved=entry) for entry in resolved_entries], - created_at=now, - updated_at=now, - ) - self._evict_completed_submissions() - self._executions[state.execution_id] = state - task = asyncio.create_task(self._submit_entries_async(state=state)) - state.task = task - self._tasks.add(task) - task.add_done_callback(self._submission_done) + async with self._admission_lock: + if self._stopping: + raise RuntimeError("Scenario suite launch scheduling is stopping.") + await self._ensure_execution_capacity_async() + + now = datetime.now(UTC) + state = _ExecutionState( + execution_id=uuid4(), + suite=suite, + entries=[_ExecutionEntryState(resolved=entry) for entry in resolved_entries], + created_at=now, + updated_at=now, + ) + self._executions[state.execution_id] = state + task = asyncio.create_task(self._submit_entries_async(state=state)) + state.task = task + self._tasks.add(task) + task.add_done_callback(self._submission_done) return await self._build_summary_async(state=state) async def get_execution_async(self, *, execution_id: UUID) -> ScenarioSuiteExecutionSummary | None: @@ -217,38 +236,20 @@ async def cancel_execution_async(self, *, execution_id: UUID) -> ScenarioSuiteEx return None current = await self._build_summary_async(state=state) - if current.status in { - ScenarioSuiteExecutionStatus.COMPLETED, - ScenarioSuiteExecutionStatus.FAILED, - ScenarioSuiteExecutionStatus.CANCELLED, - }: + if current.status in _TERMINAL_EXECUTION_STATUSES: return current - run_ids: list[str] = [] - async with state.lock: - state.cancel_requested = True - state.updated_at = datetime.now(UTC) - for entry in state.entries: - if entry.submission_task is not None and not entry.submission_task.done(): - entry.submission_task.cancel() - if entry.state == ScenarioSuiteExecutionEntryState.PENDING: - entry.state = ScenarioSuiteExecutionEntryState.CANCELLED - if entry.scenario_result_id is not None: - run_ids.append(entry.scenario_result_id) - - if run_ids: - await asyncio.gather( - *(self._get_run_service().cancel_run_async(scenario_result_id=run_id) for run_id in run_ids), - return_exceptions=True, - ) + await self._request_cancellation_async(state=state) return await self._build_summary_async(state=state) async def shutdown_async(self) -> None: - """Cancel retained submission tasks and wait for run-service cleanup.""" + """Request child cancellation and finish in-flight admission handoffs.""" self.stop_admission() + await asyncio.gather( + *(self._request_cancellation_async(state=state) for state in list(self._executions.values())), + return_exceptions=True, + ) tasks = list(self._tasks) - for task in tasks: - task.cancel() if tasks: await asyncio.gather(*tasks, return_exceptions=True) @@ -345,16 +346,14 @@ def _resolve_launch( async def _submit_entries_async(self, *, state: _ExecutionState) -> None: """Submit every resolved entry in declared order through the FIFO run service.""" try: - async with state.lock: - if state.cancel_requested: - for entry in state.entries: - entry.state = ScenarioSuiteExecutionEntryState.CANCELLED - state.updated_at = datetime.now(UTC) - return - - for entry in state.entries: + for entry in state.entries: + async with state.lock: + if state.cancel_requested: + self._mark_unstarted_entries_cancelled(state=state) + state.updated_at = datetime.now(UTC) + return entry.state = ScenarioSuiteExecutionEntryState.SUBMITTING - entry.submission_task = asyncio.create_task( + task = asyncio.create_task( self._get_run_service().start_run_async( request=entry.resolved.request, initial_metadata=entry.resolved.metadata( @@ -363,19 +362,11 @@ async def _submit_entries_async(self, *, state: _ExecutionState) -> None: ), ) ) - state.updated_at = datetime.now(UTC) + entry.submission_task = task + state.updated_at = datetime.now(UTC) - for entry in state.entries: - task = entry.submission_task - if task is None: - continue try: - result = await task - except asyncio.CancelledError: - async with state.lock: - entry.state = ScenarioSuiteExecutionEntryState.CANCELLED - state.updated_at = datetime.now(UTC) - continue + result = await asyncio.shield(task) except Exception as exc: logger.warning( "Scenario suite execution %s entry %s could not be submitted.", @@ -390,31 +381,35 @@ async def _submit_entries_async(self, *, state: _ExecutionState) -> None: state.updated_at = datetime.now(UTC) continue - async with state.lock: - entry.scenario_result_id = result.scenario_result_id - entry.state = ScenarioSuiteExecutionEntryState(result.status.value) - should_cancel = state.cancel_requested - state.updated_at = datetime.now(UTC) - if should_cancel: - await self._get_run_service().cancel_run_async(scenario_result_id=result.scenario_result_id) + await self._record_submission_result_async(state=state, entry=entry, result=result) except asyncio.CancelledError: async with state.lock: - tasks = [ - entry.submission_task + state.cancel_requested = True + self._mark_unstarted_entries_cancelled(state=state) + in_flight = [ + (entry, entry.submission_task) for entry in state.entries - if entry.submission_task is not None and not entry.submission_task.done() + if entry.state == ScenarioSuiteExecutionEntryState.SUBMITTING + and entry.submission_task is not None + and entry.scenario_result_id is None ] - for task in tasks: - task.cancel() - for entry in state.entries: - if entry.state in { - ScenarioSuiteExecutionEntryState.PENDING, - ScenarioSuiteExecutionEntryState.SUBMITTING, - }: - entry.state = ScenarioSuiteExecutionEntryState.CANCELLED state.updated_at = datetime.now(UTC) - if tasks: - await asyncio.gather(*tasks, return_exceptions=True) + for entry, task in in_flight: + try: + result = await asyncio.shield(task) + except asyncio.CancelledError: + async with state.lock: + if entry.scenario_result_id is None: + entry.state = ScenarioSuiteExecutionEntryState.CANCELLED + state.updated_at = datetime.now(UTC) + except Exception as exc: + async with state.lock: + entry.state = ScenarioSuiteExecutionEntryState.SUBMISSION_FAILED + entry.error = self.SUBMISSION_FAILURE_MESSAGE + entry.error_type = type(exc).__name__ + state.updated_at = datetime.now(UTC) + else: + await self._record_submission_result_async(state=state, entry=entry, result=result) raise async def _build_summary_async(self, *, state: _ExecutionState) -> ScenarioSuiteExecutionSummary: @@ -482,7 +477,6 @@ async def _build_summary_from_snapshot_async( entry.resolved.entry_id, exc_info=(type(run_or_error), run_or_error, run_or_error.__traceback__), ) - entry_state = ScenarioSuiteExecutionEntryState.FAILED error = self.CHILD_STATUS_FAILURE_MESSAGE error_type = type(run_or_error).__name__ elif run_or_error is not None: @@ -519,6 +513,101 @@ async def _build_summary_from_snapshot_async( entries=summaries, ) + async def _request_cancellation_async(self, *, state: _ExecutionState) -> None: + """Cancel known children without interrupting an in-flight admission handoff.""" + async with state.lock: + state.cancel_requested = True + self._mark_unstarted_entries_cancelled(state=state) + admitted_entries = [ + (entry, entry.scenario_result_id) for entry in state.entries if entry.scenario_result_id is not None + ] + state.updated_at = datetime.now(UTC) + + await asyncio.gather( + *( + self._cancel_child_run_async( + state=state, + entry=entry, + scenario_result_id=scenario_result_id, + ) + for entry, scenario_result_id in admitted_entries + ) + ) + + async def _record_submission_result_async( + self, + *, + state: _ExecutionState, + entry: _ExecutionEntryState, + result: ScenarioRunSummary, + ) -> None: + """Record an admitted child and honor cancellation requested during admission.""" + async with state.lock: + entry.scenario_result_id = result.scenario_result_id + entry.state = ScenarioSuiteExecutionEntryState(result.status.value) + entry.error = result.error + entry.error_type = result.error_type + should_cancel = state.cancel_requested + state.updated_at = datetime.now(UTC) + + if should_cancel: + await self._cancel_child_run_async( + state=state, + entry=entry, + scenario_result_id=result.scenario_result_id, + ) + + async def _cancel_child_run_async( + self, + *, + state: _ExecutionState, + entry: _ExecutionEntryState, + scenario_result_id: str, + ) -> None: + """Cancel one admitted child and retain the scheduler's authoritative state.""" + try: + result = await self._get_run_service().cancel_run_async(scenario_result_id=scenario_result_id) + except Exception: + logger.warning( + "Scenario suite execution %s child run %s could not be cancelled.", + state.execution_id, + scenario_result_id, + exc_info=True, + ) + return + if result is None: + return + + async with state.lock: + entry.state = ScenarioSuiteExecutionEntryState(result.status.value) + entry.error = result.error + entry.error_type = result.error_type + state.updated_at = datetime.now(UTC) + + @staticmethod + def _mark_unstarted_entries_cancelled(*, state: _ExecutionState) -> None: + """Mark entries that have not begun admission as cancelled.""" + for entry in state.entries: + if entry.state == ScenarioSuiteExecutionEntryState.PENDING: + entry.state = ScenarioSuiteExecutionEntryState.CANCELLED + + async def _ensure_execution_capacity_async(self) -> None: + """Evict only terminal executions or reject admission when active work fills capacity.""" + while len(self._executions) >= _MAX_RETAINED_EXECUTIONS: + evicted = False + for execution_id, state in list(self._executions.items()): + if state.task is None or not state.task.done(): + continue + summary = await self._build_summary_async(state=state) + if summary.status not in _TERMINAL_EXECUTION_STATUSES: + continue + if self._executions.get(execution_id) is state: + self._executions.pop(execution_id) + evicted = True + break + if not evicted: + raise ScenarioSuiteExecutionCapacityError(_CAPACITY_MESSAGE) + @staticmethod def _execution_status( *, @@ -557,21 +646,6 @@ def _execution_status( return ScenarioSuiteExecutionStatus.FAILED return ScenarioSuiteExecutionStatus.COMPLETED - def _evict_completed_submissions(self) -> None: - """Bound transient coordinator state after submission tasks have finished.""" - while len(self._executions) >= _MAX_RETAINED_EXECUTIONS: - execution_id = next( - ( - execution_id - for execution_id, state in self._executions.items() - if state.task is not None and state.task.done() - ), - None, - ) - if execution_id is None: - return - self._executions.pop(execution_id, None) - def _submission_done(self, task: asyncio.Task[None]) -> None: """Retrieve unexpected task failures and release active-task ownership.""" self._tasks.discard(task) diff --git a/tests/unit/backend/test_scenario_suite_execution_service.py b/tests/unit/backend/test_scenario_suite_execution_service.py index 18715f9b4d..83d0e3eaf7 100644 --- a/tests/unit/backend/test_scenario_suite_execution_service.py +++ b/tests/unit/backend/test_scenario_suite_execution_service.py @@ -13,6 +13,7 @@ import pytest +import pyrit.backend.services.scenario_suite_execution_service as _execution_mod from pyrit.backend.models.scenario_suites import ( LaunchScenarioSuiteRequest, ScenarioSuiteExecutionEntryState, @@ -21,6 +22,7 @@ ) from pyrit.backend.services.scenario_run_service import ScenarioRunService from pyrit.backend.services.scenario_suite_execution_service import ( + ScenarioSuiteExecutionCapacityError, ScenarioSuiteExecutionService, ScenarioSuiteLaunchConflictError, ScenarioSuiteNotRunnableError, @@ -422,7 +424,7 @@ async def start_run_async( assert len(calls) == 2 -async def test_child_status_failure_does_not_expose_exception_details( +async def test_child_status_failure_is_nonterminal_and_recovers( execution_context: tuple[ ScenarioSuiteExecutionService, ScenarioSuiteRegistry, @@ -432,7 +434,7 @@ async def test_child_status_failure_does_not_expose_exception_details( dict[str, ScenarioRunSummary], ], ) -> None: - """Test child status lookup failures retain diagnostics without exposing exception text.""" + """Test a temporary child status error stays nonterminal and clears after recovery.""" service, suite_registry, _, run_service, _, runs = execution_context stored = suite_registry.save_suite(suite=_suite(), expected_version=None) @@ -442,25 +444,80 @@ async def test_child_status_failure_does_not_expose_exception_details( ) await _wait_for_submissions(service) + failed_once = False + async def get_run_async(*, scenario_result_id: str) -> ScenarioRunSummary | None: - if scenario_result_id == "run-1": + nonlocal failed_once + if scenario_result_id == "run-1" and not failed_once: + failed_once = True raise RuntimeError("SELECT secret_column FROM private_table") return runs.get(scenario_result_id) runs["run-2"] = runs["run-2"].model_copy(update={"status": ScenarioRunState.COMPLETED}) run_service.get_run_async.side_effect = get_run_async - summary = await service.get_execution_async(execution_id=accepted.suite_execution_id) + unavailable = await service.get_execution_async(execution_id=accepted.suite_execution_id) + recovered = await service.get_execution_async(execution_id=accepted.suite_execution_id) + + assert unavailable is not None + assert unavailable.status == ScenarioSuiteExecutionStatus.IN_PROGRESS + assert unavailable.entries[0].state == ScenarioSuiteExecutionEntryState.QUEUED + assert unavailable.entries[0].error == service.CHILD_STATUS_FAILURE_MESSAGE + assert unavailable.entries[0].error_type == "RuntimeError" + assert "secret_column" not in unavailable.model_dump_json() + assert unavailable.entries[1].run is not None + assert recovered is not None + assert recovered.status == ScenarioSuiteExecutionStatus.IN_PROGRESS + assert recovered.entries[0].run is not None + assert recovered.entries[0].error is None + assert recovered.entries[0].error_type is None + + +async def test_child_status_failure_does_not_suppress_cancellation( + execution_context: tuple[ + ScenarioSuiteExecutionService, + ScenarioSuiteRegistry, + MagicMock, + MagicMock, + list[tuple[RunScenarioRequest, dict[str, Any]]], + dict[str, ScenarioRunSummary], + ], +) -> None: + """Test cancellation remains available while a child status read is failing.""" + service, suite_registry, _, run_service, _, runs = execution_context + stored = suite_registry.save_suite(suite=_suite(), expected_version=None) - assert summary is not None - assert summary.status == ScenarioSuiteExecutionStatus.FAILED - assert summary.entries[0].state == ScenarioSuiteExecutionEntryState.FAILED - assert summary.entries[0].error == service.CHILD_STATUS_FAILURE_MESSAGE - assert summary.entries[0].error_type == "RuntimeError" - assert "secret_column" not in summary.model_dump_json() - assert summary.entries[1].run is not None + accepted = await service.launch_async( + name="weekly", + request=_launch_request(suite_version=stored.version), + ) + await _wait_for_submissions(service) + + fail_reads = True + async def get_run_async(*, scenario_result_id: str) -> ScenarioRunSummary | None: + if scenario_result_id == "run-1" and fail_reads: + raise TimeoutError("temporary read timeout") + return runs.get(scenario_result_id) + + run_service.get_run_async.side_effect = get_run_async + cancelling = await service.cancel_execution_async(execution_id=accepted.suite_execution_id) -async def test_cancel_stops_pending_submission_and_cancels_created_children( + assert cancelling is not None + assert cancelling.status == ScenarioSuiteExecutionStatus.CANCELLED + assert run_service.cancel_run_async.await_count == 2 + + fail_reads = False + recovered = await service.get_execution_async(execution_id=accepted.suite_execution_id) + + assert recovered is not None + assert recovered.status == ScenarioSuiteExecutionStatus.CANCELLED + assert [entry.state for entry in recovered.entries] == [ + ScenarioSuiteExecutionEntryState.CANCELLED, + ScenarioSuiteExecutionEntryState.CANCELLED, + ] + + +async def test_cancel_recovers_and_cancels_child_admitted_before_response( execution_context: tuple[ ScenarioSuiteExecutionService, ScenarioSuiteRegistry, @@ -470,10 +527,11 @@ async def test_cancel_stops_pending_submission_and_cancels_created_children( dict[str, ScenarioRunSummary], ], ) -> None: - """Test best-effort whole-execution cancellation without a persisted parent.""" + """Test cancellation retains an admitted child ID until the response handoff completes.""" service, suite_registry, _, run_service, calls, runs = execution_context stored = suite_registry.save_suite(suite=_suite(), expected_version=None) - hold_second = asyncio.Event() + admitted = asyncio.Event() + release_response = asyncio.Event() async def start_run_async( *, @@ -481,32 +539,40 @@ async def start_run_async( initial_metadata: dict[str, Any], ) -> ScenarioRunSummary: calls.append((request, initial_metadata)) - if len(calls) == 2: - await hold_second.wait() summary = _run_summary( run_id=f"run-{len(calls)}", scenario_name=request.scenario_name, ) runs[summary.scenario_result_id] = summary + admitted.set() + await release_response.wait() return summary run_service.start_run_async.side_effect = start_run_async - accepted = await service.launch_async( - name="weekly", - request=_launch_request(suite_version=stored.version), - ) - for _ in range(100): - current = await service.get_execution_async(execution_id=accepted.suite_execution_id) - if current is not None and current.entries[0].run is not None: - break - await asyncio.sleep(0) - else: - raise AssertionError("First suite child was not submitted") + summary = None + try: + accepted = await service.launch_async( + name="weekly", + request=_launch_request(suite_version=stored.version), + ) + await asyncio.wait_for(admitted.wait(), timeout=5) - await service.cancel_execution_async(execution_id=accepted.suite_execution_id) - await _wait_for_submissions(service) - summary = await service.get_execution_async(execution_id=accepted.suite_execution_id) + cancelling = await service.cancel_execution_async(execution_id=accepted.suite_execution_id) + + assert cancelling is not None + assert cancelling.status == ScenarioSuiteExecutionStatus.QUEUING + assert [entry.state for entry in cancelling.entries] == [ + ScenarioSuiteExecutionEntryState.SUBMITTING, + ScenarioSuiteExecutionEntryState.CANCELLED, + ] + run_service.cancel_run_async.assert_not_awaited() + + release_response.set() + await _wait_for_submissions(service) + summary = await service.get_execution_async(execution_id=accepted.suite_execution_id) + finally: + release_response.set() assert summary is not None assert summary.status == ScenarioSuiteExecutionStatus.CANCELLED @@ -515,3 +581,42 @@ async def start_run_async( ScenarioSuiteExecutionEntryState.CANCELLED, ] run_service.cancel_run_async.assert_awaited_once_with(scenario_result_id="run-1") + assert len(calls) == 1 + + +async def test_retention_rejects_active_capacity_and_evicts_terminal_execution( + execution_context: tuple[ + ScenarioSuiteExecutionService, + ScenarioSuiteRegistry, + MagicMock, + MagicMock, + list[tuple[RunScenarioRequest, dict[str, Any]]], + dict[str, ScenarioRunSummary], + ], +) -> None: + """Test the actual retention boundary never evicts queued or running children.""" + service, suite_registry, _, _, _, runs = execution_context + stored = suite_registry.save_suite(suite=_suite(), expected_version=None) + request = _launch_request(suite_version=stored.version) + + accepted = [ + await service.launch_async(name="weekly", request=request) + for _ in range(_execution_mod._MAX_RETAINED_EXECUTIONS) + ] + await _wait_for_submissions(service) + + with pytest.raises(ScenarioSuiteExecutionCapacityError, match="still active"): + await service.launch_async(name="weekly", request=request) + + first = await service.get_execution_async(execution_id=accepted[0].suite_execution_id) + assert first is not None + assert first.status == ScenarioSuiteExecutionStatus.IN_PROGRESS + + for entry in first.entries: + assert entry.run is not None + runs[entry.run.scenario_result_id] = entry.run.model_copy(update={"status": ScenarioRunState.COMPLETED}) + + replacement = await service.launch_async(name="weekly", request=request) + + assert await service.get_execution_async(execution_id=accepted[0].suite_execution_id) is None + assert await service.get_execution_async(execution_id=replacement.suite_execution_id) is not None diff --git a/tests/unit/backend/test_scenario_suite_routes.py b/tests/unit/backend/test_scenario_suite_routes.py index c10cc75a35..568fd1152b 100644 --- a/tests/unit/backend/test_scenario_suite_routes.py +++ b/tests/unit/backend/test_scenario_suite_routes.py @@ -26,6 +26,7 @@ ScenarioSuiteResponse, ) from pyrit.backend.services.scenario_suite_execution_service import ( + ScenarioSuiteExecutionCapacityError, ScenarioSuiteExecutionNotFoundError, ScenarioSuiteExecutionService, ScenarioSuiteLaunchConflictError, @@ -247,6 +248,7 @@ def test_launch_is_rejected_during_runtime_reinitialization( (ScenarioSuiteExecutionNotFoundError("missing"), 404), (ScenarioSuiteLaunchConflictError("changed"), 409), (ScenarioSuiteNotRunnableError("not runnable"), 400), + (ScenarioSuiteExecutionCapacityError("wait for an active suite"), 429), (AzureError("storage unavailable"), 503), (RuntimeError("scheduling is stopping"), 503), ], @@ -265,6 +267,8 @@ def test_launch_maps_preflight_failures( assert response.status_code == expected_status if isinstance(error, AzureError): assert "storage unavailable" not in response.text + if isinstance(error, ScenarioSuiteExecutionCapacityError): + assert response.json()["detail"] == "wait for an active suite" def test_get_and_cancel_transient_execution( From 14c2de4fe8e3fc62e9ce63e6a58a344ab5d5d675 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Fri, 23 Oct 2026 15:36:44 -0400 Subject: [PATCH 7/7] FIX: Reset suites and resume execution polling Reset code-registered ScenarioSuites during live setup replacement while leaving stored suite documents intact. Continue polling active suite executions after consecutive transient refresh failures so stale progress can recover without a page reload. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 2403a4a9-6af6-46b2-b371-46de73a018fd --- .../hooks/useScenarioSuiteExecution.test.tsx | 24 ++++++++++-- .../src/hooks/useScenarioSuiteExecution.ts | 7 ++++ pyrit/setup/initialization.py | 2 + tests/unit/setup/test_reinitialization.py | 37 ++++++++++++++++++- 4 files changed, 65 insertions(+), 5 deletions(-) diff --git a/frontend/src/hooks/useScenarioSuiteExecution.test.tsx b/frontend/src/hooks/useScenarioSuiteExecution.test.tsx index e308c74526..5a0510e04e 100644 --- a/frontend/src/hooks/useScenarioSuiteExecution.test.tsx +++ b/frontend/src/hooks/useScenarioSuiteExecution.test.tsx @@ -110,21 +110,37 @@ describe('useScenarioSuiteExecution', () => { expect(screen.getByTestId('error')).toHaveTextContent('Execution is no longer available.') }) - it('keeps the last good execution visible when a later poll fails', async () => { + it('keeps stale progress visible and resumes polling after consecutive refresh failures', async () => { mockGetExecution .mockResolvedValueOnce(makeExecution('IN_PROGRESS')) - .mockRejectedValue(new Error('refresh failed')) + .mockRejectedValueOnce(new Error('first refresh failed')) + .mockRejectedValueOnce(new Error('second refresh failed')) + .mockResolvedValueOnce(makeExecution('COMPLETED')) render() await waitFor(() => expect(screen.getByTestId('execution-status')).toHaveTextContent('IN_PROGRESS')) await act(async () => { - await jest.advanceTimersByTimeAsync(SCENARIO_SUITE_EXECUTION_POLL_INTERVAL_MS * 2) + await jest.advanceTimersByTimeAsync(SCENARIO_SUITE_EXECUTION_POLL_INTERVAL_MS) + }) + expect(screen.getByTestId('stale')).toHaveTextContent('false') + + await act(async () => { + await jest.advanceTimersByTimeAsync(SCENARIO_SUITE_EXECUTION_POLL_INTERVAL_MS) }) expect(screen.getByTestId('execution-status')).toHaveTextContent('IN_PROGRESS') expect(screen.getByTestId('stale')).toHaveTextContent('true') - expect(screen.getByTestId('error')).toHaveTextContent('refresh failed') + expect(screen.getByTestId('error')).toHaveTextContent('second refresh failed') + + await act(async () => { + await jest.advanceTimersByTimeAsync(SCENARIO_SUITE_EXECUTION_POLL_INTERVAL_MS) + }) + + expect(screen.getByTestId('execution-status')).toHaveTextContent('COMPLETED') + expect(screen.getByTestId('stale')).toHaveTextContent('false') + expect(screen.getByTestId('error')).toHaveTextContent('') + expect(mockGetExecution).toHaveBeenCalledTimes(4) }) it('stops polling when cancellation applies a terminal response', async () => { diff --git a/frontend/src/hooks/useScenarioSuiteExecution.ts b/frontend/src/hooks/useScenarioSuiteExecution.ts index 4880878e94..c9db04f4c8 100644 --- a/frontend/src/hooks/useScenarioSuiteExecution.ts +++ b/frontend/src/hooks/useScenarioSuiteExecution.ts @@ -120,6 +120,13 @@ export function useScenarioSuiteExecution( error: apiError.detail, stale: previous !== null, }) + if ( + apiError.status !== 404 + && previous !== null + && !isTerminalSuiteExecutionStatus(previous.status) + ) { + schedulePoll() + } } } diff --git a/pyrit/setup/initialization.py b/pyrit/setup/initialization.py index dbb6fbe902..be385594f2 100644 --- a/pyrit/setup/initialization.py +++ b/pyrit/setup/initialization.py @@ -71,6 +71,7 @@ def reset_setup_registries() -> None: ConverterRegistry, InitializerRegistry, ScenarioRegistry, + ScenarioSuiteRegistry, ScorerRegistry, TargetRegistry, ) @@ -84,6 +85,7 @@ def reset_setup_registries() -> None: ConverterRegistry, InitializerRegistry, ScenarioRegistry, + ScenarioSuiteRegistry, ScorerRegistry, TargetRegistry, ): diff --git a/tests/unit/setup/test_reinitialization.py b/tests/unit/setup/test_reinitialization.py index 40530580cf..d8f6129d57 100644 --- a/tests/unit/setup/test_reinitialization.py +++ b/tests/unit/setup/test_reinitialization.py @@ -15,7 +15,15 @@ from pyrit.common.apply_defaults import get_global_default_values from pyrit.common.random_context import configure_random_seed, get_configured_random_seed from pyrit.memory import AzureSQLMemory, CentralMemory, SQLiteMemory -from pyrit.registry import AttackRegistry, InitializerRegistry, Registry, TargetRegistry +from pyrit.models.catalog import ScenarioSuite +from pyrit.registry import ( + AttackRegistry, + InitializerRegistry, + Registry, + ScenarioSuiteRegistry, + ScenarioSuiteStorage, + TargetRegistry, +) from pyrit.setup.configuration_loader import ConfigurationLoader from pyrit.setup.environment_loading import resolve_environment_async from pyrit.setup.initialization import reset_setup_registries, validate_reinitialization_memory @@ -60,6 +68,33 @@ def test_reset_setup_registries_discards_attack_catalog_without_resetting_memory assert CentralMemory.get_memory_instance() is memory +def test_reset_setup_registries_replaces_registered_suites_without_deleting_storage(tmp_path: Path) -> None: + storage = ScenarioSuiteStorage(source=str(tmp_path)) + registry = ScenarioSuiteRegistry(storage=storage) + ScenarioSuiteRegistry.set_registry_singleton(registry) + stored = registry.save_suite(suite=ScenarioSuite(name="stored_suite"), expected_version=None) + registered = ScenarioSuite(name="registered_suite") + registry.register_suite(suite=registered) + + reset_setup_registries() + + replacement = ScenarioSuiteRegistry.get_registry_singleton() + replacement.configure_source(str(tmp_path)) + assert replacement is not registry + assert replacement.get_suite(name=registered.name) is None + stored_after_reset = replacement.get_suite(name=stored.suite.name) + assert stored_after_reset is not None + assert stored_after_reset.version == stored.version + + replacement.register_suite(suite=registered) + reset_setup_registries() + + after_removal = ScenarioSuiteRegistry.get_registry_singleton() + after_removal.configure_source(str(tmp_path)) + assert after_removal.get_suite(name=registered.name) is None + assert after_removal.get_suite(name=stored.suite.name) is not None + + async def test_preflight_uses_isolated_registry_without_changing_live_state() -> None: live_registry = InitializerRegistry.get_registry_singleton() config = ConfigurationLoader(