+
+ Queued{run.queue_position ? ` · Position ${run.queue_position}` : ''}
+
+
+ {run.active_scenario_result_id
+ ? `Waiting for active run ${run.active_scenario_result_id} to finish.`
+ : 'Waiting for the scheduler to start this run.'}
+
+
+
+ Execution progress
+ Not started
+
+
+ Estimated remaining
+ Available after start
+
+
+ ) : (
@@ -404,8 +455,9 @@ function ScenarioRunPageContent({ scenarioResultId }: ScenarioRunPageContentProp
+ )}
- {isTerminalRunState(run.status) ? `Run ${formatRunState(run.status)}` : ''}
+ {queued ? progressText : isTerminalRunState(run.status) ? `Run ${formatRunState(run.status)}` : ''}
@@ -645,7 +697,9 @@ function ScenarioRunPageContent({ scenarioResultId }: ScenarioRunPageContentProp
Cancel this scenario run?
- In-flight work will be stopped. Attempts already persisted will remain available in this dashboard.
+ {queued
+ ? 'This run will be removed from the queue and will never execute.'
+ : 'In-flight work will be stopped. Attempts already persisted will remain available in this dashboard.'}
{cancelError && (
@@ -657,6 +711,7 @@ function ScenarioRunPageContent({ scenarioResultId }: ScenarioRunPageContentProp
}
onClick={() => void handleCancel()}
@@ -851,12 +906,26 @@ function objectivePreview(objective: string | null, fallbackId: string): string
if (!objective) {
return `Objective unavailable (${fallbackId})`
}
+
if (objective.length <= OBJECTIVE_PREVIEW_LENGTH) {
return objective
}
return `${objective.slice(0, OBJECTIVE_PREVIEW_LENGTH - 1)}…`
}
+function formatOverloadSummaries(
+ summaries: import('@/types').ScenarioOverloadSummary[],
+): string {
+ return summaries.map((summary) => {
+ const codes = summary.status_codes.join('/')
+ return `${formatRole(summary.component_role)} (${summary.count} × HTTP ${codes}, latest ${formatTimestamp(summary.latest_timestamp)})`
+ }).join('; ')
+}
+
+function formatRole(role: string): string {
+ return role.replace(/_/g, ' ').replace(/^\w/, (letter: string) => letter.toUpperCase())
+}
+
function shouldIgnoreAttemptRowClick(event: React.MouseEvent): boolean {
return event.button !== 0
|| hasActivationModifier(event)
diff --git a/frontend/src/hooks/useScenarioQueue.test.tsx b/frontend/src/hooks/useScenarioQueue.test.tsx
new file mode 100644
index 0000000000..f4d4ce3fae
--- /dev/null
+++ b/frontend/src/hooks/useScenarioQueue.test.tsx
@@ -0,0 +1,148 @@
+import { act, renderHook, waitFor } from '@testing-library/react'
+
+import { scenariosApi } from '@/services/api'
+import type { ScenarioQueueSnapshot } from '@/types'
+
+import { SCENARIO_QUEUE_POLL_INTERVAL_MS, useScenarioQueue } from './useScenarioQueue'
+
+jest.mock('@/services/api', () => ({
+ scenariosApi: {
+ getQueue: jest.fn(),
+ },
+}))
+
+const mockGetQueue = scenariosApi.getQueue as jest.Mock
+const FIRST_SNAPSHOT: ScenarioQueueSnapshot = {
+ revision: 1,
+ snapshot_at: '2026-01-01T00:00:00Z',
+ active: null,
+ queued: [],
+}
+const SECOND_SNAPSHOT: ScenarioQueueSnapshot = {
+ ...FIRST_SNAPSHOT,
+ revision: 2,
+ queued: [{
+ scenario_result_id: 'run-2',
+ scenario_name: 'QueuedScenario',
+ scenario_registry_name: 'queued.scenario',
+ state: 'QUEUED',
+ position: 1,
+ created_at: '2026-01-01T00:00:01Z',
+ enqueued_at: '2026-01-01T00:00:01Z',
+ }],
+}
+
+describe('useScenarioQueue', () => {
+ beforeEach(() => {
+ jest.clearAllMocks()
+ jest.useFakeTimers()
+ })
+
+ afterEach(() => {
+ jest.useRealTimers()
+ })
+
+ it('polls and applies position changes', async () => {
+ mockGetQueue
+ .mockResolvedValueOnce(FIRST_SNAPSHOT)
+ .mockResolvedValueOnce(SECOND_SNAPSHOT)
+
+ const { result, unmount } = renderHook(() => useScenarioQueue())
+ await waitFor(() => expect(result.current.snapshot?.revision).toBe(1))
+ await act(async () => {
+ await jest.advanceTimersByTimeAsync(SCENARIO_QUEUE_POLL_INTERVAL_MS)
+ })
+
+ expect(result.current.snapshot?.queued[0].position).toBe(1)
+ unmount()
+ })
+
+ it('keeps the last good snapshot across a transient failure', async () => {
+ mockGetQueue
+ .mockResolvedValueOnce(FIRST_SNAPSHOT)
+ .mockRejectedValueOnce(new Error('temporary queue failure'))
+
+ const { result, unmount } = renderHook(() => useScenarioQueue())
+ await waitFor(() => expect(result.current.snapshot?.revision).toBe(1))
+ await act(async () => {
+ await jest.advanceTimersByTimeAsync(SCENARIO_QUEUE_POLL_INTERVAL_MS)
+ })
+
+ expect(result.current.snapshot).toEqual(FIRST_SNAPSHOT)
+ expect(result.current.stale).toBe(true)
+ expect(result.current.error).toBe('temporary queue failure')
+ unmount()
+ })
+
+ it('retries an initial failure without presenting stale queue data', async () => {
+ mockGetQueue
+ .mockRejectedValueOnce(new Error('queue unavailable'))
+ .mockResolvedValueOnce(FIRST_SNAPSHOT)
+
+ const { result, unmount } = renderHook(() => useScenarioQueue())
+ await waitFor(() => expect(result.current.error).toBe('queue unavailable'))
+
+ expect(result.current.snapshot).toBeNull()
+ expect(result.current.loading).toBe(false)
+ expect(result.current.stale).toBe(false)
+
+ act(() => result.current.retry())
+ expect(result.current.loading).toBe(true)
+ expect(result.current.error).toBeNull()
+ await waitFor(() => expect(result.current.snapshot).toEqual(FIRST_SNAPSHOT))
+ unmount()
+ })
+
+ it('keeps existing queue data visible while a manual retry is pending', async () => {
+ let resolveRetry: ((snapshot: ScenarioQueueSnapshot) => void) | undefined
+ mockGetQueue
+ .mockResolvedValueOnce(FIRST_SNAPSHOT)
+ .mockImplementationOnce(() => new Promise((resolve) => {
+ resolveRetry = resolve
+ }))
+
+ const { result, unmount } = renderHook(() => useScenarioQueue())
+ await waitFor(() => expect(result.current.snapshot).toEqual(FIRST_SNAPSHOT))
+
+ act(() => result.current.retry())
+
+ expect(result.current.loading).toBe(false)
+ expect(result.current.snapshot).toEqual(FIRST_SNAPSHOT)
+ await act(async () => {
+ resolveRetry?.(SECOND_SNAPSHOT)
+ })
+ unmount()
+ })
+
+ it('ignores a request that resolves after unmount', async () => {
+ let resolveRequest: ((snapshot: ScenarioQueueSnapshot) => void) | undefined
+ mockGetQueue.mockImplementationOnce(() => new Promise((resolve) => {
+ resolveRequest = resolve
+ }))
+
+ const { result, unmount } = renderHook(() => useScenarioQueue())
+ await waitFor(() => expect(mockGetQueue).toHaveBeenCalledTimes(1))
+ unmount()
+
+ await act(async () => {
+ resolveRequest?.(FIRST_SNAPSHOT)
+ })
+ expect(result.current.snapshot).toBeNull()
+ })
+
+ it('ignores a request that rejects after unmount', async () => {
+ let rejectRequest: ((reason?: unknown) => void) | undefined
+ mockGetQueue.mockImplementationOnce(() => new Promise((_resolve, reject) => {
+ rejectRequest = reject
+ }))
+
+ const { result, unmount } = renderHook(() => useScenarioQueue())
+ await waitFor(() => expect(mockGetQueue).toHaveBeenCalledTimes(1))
+ unmount()
+
+ await act(async () => {
+ rejectRequest?.(new Error('late failure'))
+ })
+ expect(result.current.error).toBeNull()
+ })
+})
diff --git a/frontend/src/hooks/useScenarioQueue.ts b/frontend/src/hooks/useScenarioQueue.ts
new file mode 100644
index 0000000000..3586ec3128
--- /dev/null
+++ b/frontend/src/hooks/useScenarioQueue.ts
@@ -0,0 +1,82 @@
+import { useCallback, useEffect, useRef, useState } from 'react'
+
+import { scenariosApi } from '@/services/api'
+import { toApiError } from '@/services/errors'
+import type { ScenarioQueueSnapshot } from '@/types'
+
+export const SCENARIO_QUEUE_POLL_INTERVAL_MS = 2_500
+
+export interface ScenarioQueueState {
+ readonly snapshot: ScenarioQueueSnapshot | null
+ readonly loading: boolean
+ readonly stale: boolean
+ readonly error: string | null
+}
+
+export interface UseScenarioQueueResult extends ScenarioQueueState {
+ readonly retry: () => void
+}
+
+export function useScenarioQueue(): UseScenarioQueueResult {
+ const [state, setState] = useState({
+ snapshot: null,
+ loading: true,
+ stale: false,
+ error: null,
+ })
+ const [retryEpoch, setRetryEpoch] = useState(0)
+ const timerRef = useRef | null>(null)
+
+ useEffect(() => {
+ let active = true
+ const controller = new AbortController()
+
+ const fetchQueueAsync = async (): Promise => {
+ if (!active) {
+ return
+ }
+ try {
+ const snapshot = await scenariosApi.getQueue(controller.signal)
+ if (!active) {
+ return
+ }
+ setState({ snapshot, loading: false, stale: false, error: null })
+ } catch (error: unknown) {
+ if (!active || controller.signal.aborted) {
+ return
+ }
+ const message = toApiError(error).detail
+ setState((previous) => ({
+ ...previous,
+ loading: false,
+ stale: previous.snapshot !== null,
+ error: message,
+ }))
+ } finally {
+ if (active) {
+ timerRef.current = setTimeout(() => {
+ timerRef.current = null
+ void fetchQueueAsync()
+ }, SCENARIO_QUEUE_POLL_INTERVAL_MS)
+ }
+ }
+ }
+
+ void fetchQueueAsync()
+ return () => {
+ active = false
+ controller.abort()
+ if (timerRef.current !== null) {
+ clearTimeout(timerRef.current)
+ timerRef.current = null
+ }
+ }
+ }, [retryEpoch])
+
+ const retry = useCallback((): void => {
+ setState((previous) => ({ ...previous, loading: previous.snapshot === null, stale: false, error: null }))
+ setRetryEpoch((epoch) => epoch + 1)
+ }, [])
+
+ return { ...state, retry }
+}
diff --git a/frontend/src/services/api.ts b/frontend/src/services/api.ts
index 4d49dad623..75580637d0 100644
--- a/frontend/src/services/api.ts
+++ b/frontend/src/services/api.ts
@@ -36,6 +36,7 @@ import type {
ScenarioRunSummary,
ScenarioRunListResponse,
ScenarioRunProgress,
+ ScenarioQueueSnapshot,
ScenarioRunState,
} from '../types'
@@ -423,6 +424,11 @@ export const scenariosApi = {
return response.data
},
+ getQueue: async (signal?: AbortSignal): Promise => {
+ const response = await apiClient.get('/scenarios/runs/queue', { signal })
+ return response.data
+ },
+
cancelRun: async (scenarioResultId: string, signal?: AbortSignal): Promise => {
const response = await apiClient.post(
`/scenarios/runs/${encodeURIComponent(scenarioResultId)}/cancel`,
diff --git a/frontend/src/types/index.ts b/frontend/src/types/index.ts
index 026842fc69..82d9bfe4bf 100644
--- a/frontend/src/types/index.ts
+++ b/frontend/src/types/index.ts
@@ -605,6 +605,7 @@ export interface RetryEvent {
component_role: string
component_name?: string | null
endpoint?: string | null
+ status_code?: number | null
elapsed_seconds: number
}
@@ -616,6 +617,15 @@ export interface AttackRetrySummary {
export type ScenarioRunState = 'CREATED' | 'QUEUED' | 'IN_PROGRESS' | 'COMPLETED' | 'FAILED' | 'CANCELLED'
+export interface ScenarioOverloadSummary {
+ component_role: string
+ count: number
+ rate_limit_count: number
+ server_error_count: number
+ status_codes: number[]
+ latest_timestamp: string
+}
+
export interface ScenarioRunSummary {
scenario_result_id: string
scenario_name: string
@@ -623,6 +633,7 @@ export interface ScenarioRunSummary {
scenario_version: number
status: ScenarioRunState
created_at: string
+ started_at?: string | null
updated_at: string
error?: string | null
error_type?: string | null
@@ -643,6 +654,9 @@ export interface ScenarioRunSummary {
successful_attacks?: number
error_attacks?: number
attack_details_available?: boolean
+ queue_position?: number | null
+ active_scenario_result_id?: string | null
+ overload_summaries?: ScenarioOverloadSummary[]
}
export interface ScenarioTargetSummary {
@@ -665,6 +679,7 @@ export interface ScenarioProgressHeader {
scenario_version: number
status: ScenarioRunState
created_at: string
+ started_at?: string | null
completed_at?: string | null
pyrit_version?: string | null
target?: ScenarioTargetSummary | null
@@ -672,6 +687,27 @@ export interface ScenarioProgressHeader {
datasets_used?: string[]
scenario_parameters?: Record
labels?: Record
+ queue_position?: number | null
+ active_scenario_result_id?: string | null
+ overload_summaries?: ScenarioOverloadSummary[]
+}
+
+export interface ScenarioQueueEntry {
+ scenario_result_id: string
+ scenario_name: string
+ scenario_registry_name: string
+ created_at: string
+ enqueued_at: string
+ started_at?: string | null
+ state: ScenarioRunState
+ position?: number | null
+}
+
+export interface ScenarioQueueSnapshot {
+ revision: number
+ snapshot_at: string
+ active?: ScenarioQueueEntry | null
+ queued: ScenarioQueueEntry[]
}
/** One persisted attack attempt in ascending progress order. */
diff --git a/frontend/src/utils/scenarioRunProgress.test.ts b/frontend/src/utils/scenarioRunProgress.test.ts
index 414b35bd17..b9ac1fbacd 100644
--- a/frontend/src/utils/scenarioRunProgress.test.ts
+++ b/frontend/src/utils/scenarioRunProgress.test.ts
@@ -72,6 +72,7 @@ function makePage(overrides: Partial = {}): ScenarioRunProg
scenario_version: 1,
status: 'IN_PROGRESS',
created_at: '2026-01-01T00:00:00Z',
+ started_at: '2026-01-01T00:00:00Z',
},
plan: PLAN,
reset: false,
@@ -119,6 +120,63 @@ describe('scenarioRunProgressReducer', () => {
expect(reset.cursor).toBe('cursor-2')
})
+ it('does not double-count overload evidence during cancellation catch-up', () => {
+ const overload = {
+ component_role: 'objective_target',
+ count: 1,
+ rate_limit_count: 1,
+ server_error_count: 0,
+ status_codes: [429],
+ latest_timestamp: '2026-01-01T00:01:00Z',
+ }
+ const first = scenarioRunProgressReducer(INITIAL_SCENARIO_RUN_PROGRESS_STATE, {
+ type: 'apply-page',
+ page: makePage({ run: { ...makePage().run, overload_summaries: [overload] } }),
+ fresh: true,
+ })
+ const cancelled = scenarioRunProgressReducer(first, {
+ type: 'apply-run-summary',
+ run: {
+ scenario_result_id: 'run-1',
+ scenario_name: 'TestScenario',
+ scenario_registry_name: 'test.scenario',
+ scenario_version: 1,
+ status: 'CANCELLED',
+ created_at: '2026-01-01T00:00:00Z',
+ started_at: '2026-01-01T00:00:30Z',
+ updated_at: '2026-01-01T00:02:00Z',
+ techniques_used: [],
+ total_attacks: 2,
+ completed_attacks: 2,
+ objective_achieved_rate: 0,
+ failed_attacks: [],
+ attack_retries: [],
+ total_retries: 2,
+ labels: {},
+ overload_summaries: [{ ...overload, count: 2, rate_limit_count: 2 }],
+ },
+ })
+ const caughtUp = scenarioRunProgressReducer(cancelled, {
+ type: 'apply-page',
+ page: makePage({
+ plan: null,
+ run: {
+ ...makePage().run,
+ status: 'CANCELLED',
+ overload_summaries: [{
+ ...overload,
+ latest_timestamp: '2026-01-01T00:02:00Z',
+ }],
+ },
+ }),
+ fresh: false,
+ })
+
+ expect(cancelled.overloadSummaries[0].count).toBe(1)
+ expect(cancelled.run?.started_at).toBe('2026-01-01T00:00:30Z')
+ expect(caughtUp.overloadSummaries[0].count).toBe(2)
+ })
+
it('retains last-good data and marks it stale after a transient failure', () => {
const first = readyState([makeResult('attempt-1', 'group-a', 'seed-1', 'success', 1)])
const failed = scenarioRunProgressReducer(first, {
@@ -240,16 +298,45 @@ describe('scenario run progress calculations', () => {
])
})
- it('uses now for active elapsed time and completed_at for terminal elapsed time', () => {
- const active = makePage().run
- expect(getElapsedMilliseconds(active, Date.parse('2026-01-01T00:05:00Z'))).toBe(300_000)
+ it('uses execution start for elapsed time and excludes a long queue delay from active ETA', () => {
+ const active = {
+ ...makePage().run,
+ created_at: '2026-01-01T00:00:00Z',
+ started_at: '2026-01-01T01:00:00Z',
+ }
+ expect(getElapsedMilliseconds(active, Date.parse('2026-01-01T01:05:00Z'))).toBe(300_000)
const terminal = {
...active,
status: 'COMPLETED' as const,
- completed_at: '2026-01-01T00:03:00Z',
+ completed_at: '2026-01-01T01:03:00Z',
}
- expect(getElapsedMilliseconds(terminal, Date.parse('2026-01-01T00:05:00Z'))).toBe(180_000)
+ expect(getElapsedMilliseconds(terminal, Date.parse('2026-01-01T01:05:00Z'))).toBe(180_000)
+
+ const state = readyState([makeResult('a-1', 'group-a', 'seed-1', 'success', 1)])
+ state.run = active
+ expect(getEtaMilliseconds(state, Date.parse('2026-01-01T01:02:00Z'))).toBe(240_000)
+ })
+
+ it('does not count queue wait as elapsed time or fabricate a queued ETA', () => {
+ const queued = {
+ ...makePage().run,
+ status: 'QUEUED' as const,
+ created_at: '2026-01-01T00:00:00Z',
+ started_at: null,
+ }
+ const now = Date.parse('2026-01-01T01:00:00Z')
+
+ expect(getElapsedMilliseconds(queued, now)).toBe(0)
+
+ const state = readyState([makeResult('a-1', 'group-a', 'seed-1', 'success', 1)])
+ state.run = queued
+ expect(getEtaMilliseconds(state, now)).toBeNull()
+
+ expect(getElapsedMilliseconds(
+ { ...queued, status: 'CANCELLED', completed_at: '2026-01-01T00:30:00Z' },
+ now,
+ )).toBe(0)
})
it('calculates ETA from observed wall-clock completion rate and hides unsafe estimates', () => {
diff --git a/frontend/src/utils/scenarioRunProgress.ts b/frontend/src/utils/scenarioRunProgress.ts
index bac819f9c6..c442dc6cd9 100644
--- a/frontend/src/utils/scenarioRunProgress.ts
+++ b/frontend/src/utils/scenarioRunProgress.ts
@@ -1,6 +1,7 @@
import type {
ScenarioProgressHeader,
ScenarioProgressResult,
+ ScenarioOverloadSummary,
ScenarioRunPlan,
ScenarioRunPlanAtomicGroup,
ScenarioRunState,
@@ -21,6 +22,7 @@ export interface ScenarioRunProgressState {
readonly hasMore: boolean
readonly error: string | null
readonly stale: boolean
+ readonly overloadSummaries: ScenarioOverloadSummary[]
}
export type ScenarioRunProgressAction =
@@ -90,6 +92,7 @@ export const INITIAL_SCENARIO_RUN_PROGRESS_STATE: ScenarioRunProgressState = {
hasMore: false,
error: null,
stale: false,
+ overloadSummaries: [],
}
export function isTerminalRunState(status: ScenarioRunState): boolean {
@@ -131,6 +134,7 @@ export function scenarioRunProgressReducer(
scenario_version: action.run.scenario_version,
status: action.run.status,
created_at: action.run.created_at,
+ started_at: action.run.started_at,
completed_at: action.run.completed_at,
pyrit_version: action.run.pyrit_version,
target: action.run.target,
@@ -138,11 +142,15 @@ export function scenarioRunProgressReducer(
datasets_used: action.run.datasets_used ?? [],
scenario_parameters: action.run.scenario_parameters ?? {},
labels: action.run.labels,
+ queue_position: action.run.queue_position,
+ active_scenario_result_id: action.run.active_scenario_result_id,
+ overload_summaries: action.run.overload_summaries ?? [],
},
activeAtomicGroupIds: [],
error: null,
stale: false,
hasMore: false,
+ overloadSummaries: state.overloadSummaries,
}
}
@@ -158,6 +166,10 @@ export function scenarioRunProgressReducer(
}
const results = [...resultsById.values()].sort(compareAttempts)
+ const overloadSummaries = mergeOverloadSummaries(
+ shouldReset ? [] : state.overloadSummaries,
+ action.page.run.overload_summaries ?? [],
+ )
return {
loadStatus: 'ready',
run: action.page.run,
@@ -169,9 +181,33 @@ export function scenarioRunProgressReducer(
hasMore: action.page.has_more,
error: null,
stale: false,
+ overloadSummaries,
}
}
+function mergeOverloadSummaries(
+ existing: ScenarioOverloadSummary[],
+ incoming: ScenarioOverloadSummary[],
+): ScenarioOverloadSummary[] {
+ const byRole = new Map(existing.map((summary) => [summary.component_role, summary]))
+ for (const summary of incoming) {
+ const previous = byRole.get(summary.component_role)
+ byRole.set(summary.component_role, previous ? {
+ component_role: summary.component_role,
+ count: previous.count + summary.count,
+ rate_limit_count: previous.rate_limit_count + summary.rate_limit_count,
+ server_error_count: previous.server_error_count + summary.server_error_count,
+ status_codes: [...new Set([...previous.status_codes, ...summary.status_codes])].sort((left, right) => left - right),
+ latest_timestamp: Date.parse(summary.latest_timestamp) >= Date.parse(previous.latest_timestamp)
+ ? summary.latest_timestamp
+ : previous.latest_timestamp,
+ } : summary)
+ }
+ return [...byRole.values()].sort(
+ (left, right) => Date.parse(right.latest_timestamp) - Date.parse(left.latest_timestamp),
+ )
+}
+
export function getOverallProgress(state: ScenarioRunProgressState): OverallProgress {
const units = buildUnitAttempts(state.results)
const completed = [...units.values()].filter((unit) => unit.latestNonError !== null).length
@@ -198,15 +234,18 @@ export function getElapsedMilliseconds(
run: ScenarioProgressHeader,
nowMilliseconds: number,
): number {
- const created = Date.parse(run.created_at)
+ if (!run.started_at) {
+ return 0
+ }
+ const started = Date.parse(run.started_at)
const terminalEnd = run.completed_at ? Date.parse(run.completed_at) : Number.NaN
const end = isTerminalRunState(run.status) && Number.isFinite(terminalEnd)
? terminalEnd
: nowMilliseconds
- if (!Number.isFinite(created) || !Number.isFinite(end)) {
+ if (!Number.isFinite(started) || !Number.isFinite(end)) {
return 0
}
- return Math.max(0, end - created)
+ return Math.max(0, end - started)
}
export function getEtaMilliseconds(
diff --git a/pyrit/backend/main.py b/pyrit/backend/main.py
index 18825fa3ac..c3fda214a7 100644
--- a/pyrit/backend/main.py
+++ b/pyrit/backend/main.py
@@ -35,7 +35,7 @@
targets,
version,
)
-from pyrit.backend.services.initializer_service import get_initializer_service
+from pyrit.backend.services import get_initializer_service, get_scenario_run_service
from pyrit.setup.configuration_loader import ConfigurationLoader
# Check for development mode from environment variable
@@ -84,12 +84,18 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
if config.allow_custom_initializers:
logger.warning("Custom initializer registration is ENABLED (allow_custom_initializers: true).")
+ scenario_run_service = get_scenario_run_service()
+ await scenario_run_service.reconcile_interrupted_runs_async()
+
# Mount the bundled frontend (or print a dev/missing-frontend notice).
# Done here rather than at module load so test imports of `pyrit.backend.main`
# don't emit noise and don't perform filesystem side effects.
setup_frontend()
- yield
+ try:
+ yield
+ finally:
+ await scenario_run_service.shutdown_async()
app = FastAPI(
diff --git a/pyrit/backend/routes/scenarios.py b/pyrit/backend/routes/scenarios.py
index 8197866a53..b43e643ce5 100644
--- a/pyrit/backend/routes/scenarios.py
+++ b/pyrit/backend/routes/scenarios.py
@@ -22,15 +22,14 @@
)
from pyrit.backend.services.scenario_run_service import get_scenario_run_service
from pyrit.backend.services.scenario_service import get_scenario_service
-from pyrit.models import ScenarioResult, ScenarioRunState
-from pyrit.models.catalog.scenario import (
+from pyrit.models import ScenarioQueueSnapshot, ScenarioResult, ScenarioRunProgress, ScenarioRunState
+from pyrit.models.catalog import (
RegisteredScenario,
RunScenarioRequest,
ScenarioDefaultRunSizeEstimate,
ScenarioRunSizeEstimateRequest,
ScenarioRunSummary,
)
-from pyrit.models.scenario_progress import ScenarioRunProgress
router = APIRouter(prefix="/scenarios", tags=["scenarios"])
@@ -228,6 +227,20 @@ async def list_scenario_runs( # pyrit-async-suffix-exempt
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None
+@router.get(
+ "/runs/queue",
+ response_model=ScenarioQueueSnapshot,
+)
+async def get_scenario_run_queue() -> ScenarioQueueSnapshot: # pyrit-async-suffix-exempt
+ """
+ Get the active scenario and ordered FIFO waiting queue.
+
+ Returns:
+ ScenarioQueueSnapshot: Current in-process scheduler state.
+ """
+ return get_scenario_run_service().get_queue_snapshot()
+
+
@router.get(
"/runs/{scenario_result_id}",
response_model=ScenarioRunSummary,
@@ -251,6 +264,8 @@ async def get_scenario_run(scenario_result_id: str) -> ScenarioRunSummary: # py
service.get_run_from_storage,
scenario_result_id=scenario_result_id,
active_error=active_snapshot.error,
+ queue_position=active_snapshot.queue_position,
+ active_scenario_result_id=active_snapshot.active_scenario_result_id,
)
if run is None:
raise HTTPException(
@@ -289,6 +304,8 @@ async def get_scenario_run_progress( # pyrit-async-suffix-exempt
since=since,
limit=limit,
active_group_ids=active_snapshot.active_group_ids,
+ queue_position=active_snapshot.queue_position,
+ active_scenario_result_id=active_snapshot.active_scenario_result_id,
)
except ValueError as exc:
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None
diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py
index 763a204e3f..d407d49a10 100644
--- a/pyrit/backend/services/scenario_run_service.py
+++ b/pyrit/backend/services/scenario_run_service.py
@@ -16,18 +16,24 @@
import json
import logging
import uuid
+from collections import OrderedDict, deque
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any
from urllib.parse import urlsplit, urlunsplit
+try:
+ from builtins import ExceptionGroup # type: ignore[attr-defined,ty:unresolved-import]
+except ImportError: # pragma: no cover - exercised only on 3.10
+ from exceptiongroup import ExceptionGroup # type: ignore[no-redef,ty:unresolved-import]
+
from pydantic import TypeAdapter, ValidationError
from pyrit.backend.models.common import PaginationInfo, filter_sensitive_fields
from pyrit.backend.models.scenarios import ScenarioRunListResponse
from pyrit.common.utils import to_sha256
-from pyrit.memory import CentralMemory
+from pyrit.memory import CentralMemory, SQLiteMemory
from pyrit.memory.memory_interface import (
ScenarioHistoryKeysetCursor,
ScenarioHistoryRunRecord,
@@ -36,6 +42,7 @@
)
from pyrit.models import (
SCENARIO_RUN_PLAN_METADATA_KEY,
+ SCENARIO_RUN_STARTED_AT_METADATA_KEY,
AtomicAttackIdentifier,
AttackOutcome,
ComponentIdentifier,
@@ -43,6 +50,8 @@
ScenarioIdentifier,
ScenarioProgressHeader,
ScenarioProgressResult,
+ ScenarioQueueEntry,
+ ScenarioQueueSnapshot,
ScenarioResult,
ScenarioRunPlan,
ScenarioRunPlanAtomicGroup,
@@ -56,6 +65,7 @@
AttackErrorSummary,
AttackRetrySummary,
RunScenarioRequest,
+ ScenarioOverloadSummary,
ScenarioRunSummary,
ScenarioTargetSummary,
)
@@ -73,7 +83,21 @@
logger = logging.getLogger(__name__)
-_DEFAULT_MAX_CONCURRENT_RUNS = 3
+_DEFAULT_MAX_CONCURRENT_RUNS = 1
+_MAX_OVERLOAD_EVENTS = 500
+_MAX_OVERLOAD_ROLES = 16
+_MAX_TERMINAL_ERRORS = 100
+_SCHEDULER_RETRY_INITIAL_SECONDS = 0.05
+_SCHEDULER_RETRY_MAX_SECONDS = 1.0
+_SCHEDULER_METADATA_KEY = "scheduler_managed_by"
+_SCHEDULER_METADATA_VALUE = "ScenarioRunService.process_local_fifo"
+_INTERRUPTED_ERROR_TYPE = "ScenarioInterruptedError"
+_RESTART_INTERRUPTION_REASON = (
+ "The backend process restarted before this scenario run completed; "
+ "its executable scenario objects could not be recovered safely."
+)
+_SHUTDOWN_INTERRUPTION_REASON = "The backend process shut down before this scenario run completed."
+_USER_CANCELLATION_REASON = "Run was cancelled by user"
_CONVERTER_MODIFIER_PREFIX = "converter."
_SAFE_SCENARIO_PARAMETER_NAMES = frozenset(
@@ -90,6 +114,7 @@
)
_HISTORY_ATOMIC_GROUPS_ADAPTER = TypeAdapter(list[ScenarioRunPlanAtomicGroup])
_HISTORY_SEED_ID_MAP_ADAPTER = TypeAdapter(list[dict[str, str]])
+_STARTED_AT_ADAPTER = TypeAdapter(datetime)
@dataclass
@@ -100,6 +125,15 @@ class _ActiveTask:
task: asyncio.Task[None] | None = None
scenario: Scenario | None = None
error: str | None = None
+ scenario_name: str = ""
+ scenario_registry_name: str = ""
+ created_at: datetime | None = None
+ enqueued_at: datetime | None = None
+ started_at: datetime | None = None
+ cancellation_state: ScenarioRunState = ScenarioRunState.CANCELLED
+ cancellation_reason: str = _USER_CANCELLATION_REASON
+ cancellation_error_type: str = "CancelledError"
+ retain_error_on_terminalization: bool = False
@dataclass(frozen=True, slots=True)
@@ -108,6 +142,8 @@ class _ActiveRunSnapshot:
error: str | None = None
active_group_ids: tuple[str, ...] = ()
+ queue_position: int | None = None
+ active_scenario_result_id: str | None = None
class ScenarioRunService:
@@ -115,45 +151,50 @@ class ScenarioRunService:
Service for managing scenario run lifecycle.
Uses CentralMemory (database) as the source of truth for run state.
- Keeps an in-memory dict only for active asyncio tasks (cancellation support).
+ Keeps executable objects in a process-local single-active FIFO scheduler.
"""
def __init__(self, *, max_concurrent_runs: int = _DEFAULT_MAX_CONCURRENT_RUNS) -> None:
- """Initialize the scenario run service."""
- self._max_concurrent_runs = max_concurrent_runs
+ """
+ Initialize the scenario run service.
+
+ ``max_concurrent_runs`` remains accepted for configuration compatibility;
+ scenario execution is always serialized to one active run.
+ """
+ if max_concurrent_runs < 1:
+ raise ValueError("max_concurrent_runs must be at least 1.")
self._memory = CentralMemory.get_memory_instance()
self._active_tasks: dict[str, _ActiveTask] = {}
- self._run_semaphore = asyncio.Semaphore(max_concurrent_runs)
+ self._terminal_errors: OrderedDict[str, str] = OrderedDict()
+ self._active_scenario_result_id: str | None = None
+ self._queued_runs: deque[_ActiveTask] = deque()
+ self._handoff_retry_tasks: set[asyncio.Task[None]] = set()
+ self._scheduler_lock = asyncio.Lock()
+ self._launch_lock = asyncio.Lock()
+ self._queue_revision = 0
+ self._stopping = False
async def start_run_async(self, *, request: RunScenarioRequest) -> ScenarioRunSummary:
"""
- Start a new scenario run as a background task.
+ Initialize and schedule a scenario run.
Performs all validation and initialization eagerly (initializers, target
resolution, technique validation, scenario.initialize_async) so errors are
- returned immediately. On success, spawns a background task that only
- executes scenario.run_async.
+ returned immediately. On success, starts execution when idle or appends
+ the initialized run to the FIFO waiting queue.
Args:
request: The run request with scenario name, target, and options.
Returns:
- ScenarioRunResponse with run_id and RUNNING status.
+ ScenarioRunSummary with a stable ID and current active or queued state.
Raises:
- ValueError: If scenario, target, initializer, or technique cannot be found,
- or concurrent limit exceeded.
+ ValueError: If scenario, target, initializer, or technique cannot be found.
"""
- if self._run_semaphore.locked():
- raise ValueError(
- f"Maximum concurrent runs ({self._max_concurrent_runs}) reached. "
- "Wait for an existing run to complete or cancel one."
- )
-
- await self._run_semaphore.acquire()
-
- # Perform all initialization eagerly — errors propagate to caller
- try:
+ async with self._launch_lock:
+ if self._stopping:
+ raise RuntimeError("Scenario run scheduling is stopping.")
scenario_class = self._resolve_scenario_class(request=request)
await self._run_initializers_async(request=request)
objective_target = self._resolve_target(request=request)
@@ -161,24 +202,33 @@ async def start_run_async(self, *, request: RunScenarioRequest) -> ScenarioRunSu
request=request, scenario_class=scenario_class, objective_target=objective_target
)
scenario = await self._initialize_scenario_async(request=request, init_kwargs=init_kwargs)
- except Exception:
- self._run_semaphore.release()
- raise
-
- # scenario_result_id is set during initialize_async
- scenario_result_id = scenario._scenario_result_id
- if scenario_result_id is None:
- raise ValueError("Scenario did not produce a scenario_result_id during initialization.")
-
- # Track active task
- active = _ActiveTask(scenario_result_id=scenario_result_id, scenario=scenario)
- self._active_tasks[scenario_result_id] = active
-
- # Spawn background task (only runs scenario.run_async)
- task = asyncio.create_task(self._execute_run_async(scenario_result_id=scenario_result_id))
- active.task = task
+ scenario_result_id = scenario._scenario_result_id
+ if scenario_result_id is None:
+ raise ValueError("Scenario did not produce a scenario_result_id during initialization.")
+ persisted = await asyncio.to_thread(
+ self._memory.get_scenario_results,
+ scenario_result_ids=[scenario_result_id],
+ )
+ if not persisted:
+ raise RuntimeError(f"Scenario run {scenario_result_id} was not persisted during initialization.")
+ scheduled = _ActiveTask(
+ scenario_result_id=scenario_result_id,
+ scenario=scenario,
+ scenario_name=persisted[0].scenario_name,
+ scenario_registry_name=request.scenario_name,
+ created_at=persisted[0].creation_time,
+ enqueued_at=datetime.now(timezone.utc),
+ )
+ await self._enqueue_run_async(scheduled=scheduled)
- response = self.get_run(scenario_result_id=scenario_result_id)
+ snapshot = self.snapshot_active_run(scenario_result_id=scenario_result_id)
+ response = await asyncio.to_thread(
+ self.get_run_from_storage,
+ scenario_result_id=scenario_result_id,
+ active_error=snapshot.error,
+ queue_position=snapshot.queue_position,
+ active_scenario_result_id=snapshot.active_scenario_result_id,
+ )
if response is None:
raise RuntimeError(f"Scenario run {scenario_result_id} was not found in the database after initialization.")
return response
@@ -194,13 +244,20 @@ def get_run(self, *, scenario_result_id: str) -> ScenarioRunSummary | None:
ScenarioRunSummary if found, None otherwise.
"""
snapshot = self.snapshot_active_run(scenario_result_id=scenario_result_id)
- return self.get_run_from_storage(scenario_result_id=scenario_result_id, active_error=snapshot.error)
+ return self.get_run_from_storage(
+ scenario_result_id=scenario_result_id,
+ active_error=snapshot.error,
+ queue_position=snapshot.queue_position,
+ active_scenario_result_id=snapshot.active_scenario_result_id,
+ )
def get_run_from_storage(
self,
*,
scenario_result_id: str,
active_error: str | None,
+ queue_position: int | None = None,
+ active_scenario_result_id: str | None = None,
) -> ScenarioRunSummary | None:
"""
Build a run summary using database state plus an event-loop snapshot.
@@ -208,11 +265,18 @@ def get_run_from_storage(
Args:
scenario_result_id: The scenario result ID.
active_error: Error copied from the active asyncio task, if any.
+ queue_position: Current 1-based waiting position, if queued.
+ active_scenario_result_id: Currently executing scenario result ID.
Returns:
ScenarioRunSummary | None: The run summary when found.
"""
- return self._build_response(scenario_result_id=scenario_result_id, active_error=active_error)
+ return self._build_response(
+ scenario_result_id=scenario_result_id,
+ active_error=active_error,
+ queue_position=queue_position,
+ active_scenario_result_id=active_scenario_result_id,
+ )
def list_runs(
self,
@@ -299,33 +363,387 @@ async def cancel_run_async(self, *, scenario_result_id: str) -> ScenarioRunSumma
Raises:
ValueError: If the run is already in a terminal state or not active.
"""
- # Verify run exists in DB
- results = self._memory.get_scenario_results(scenario_result_ids=[scenario_result_id])
+ results = await asyncio.to_thread(
+ self._memory.get_scenario_results,
+ scenario_result_ids=[scenario_result_id],
+ )
if not results:
return None
- scenario_result = results[0]
- db_status = scenario_result.scenario_run_state
-
- if db_status in (ScenarioRunState.COMPLETED, ScenarioRunState.FAILED, ScenarioRunState.CANCELLED):
+ db_status = results[0].scenario_run_state
+ if self._is_terminal_state(db_status):
raise ValueError(f"Cannot cancel run in '{db_status}' state.")
- # Cancel the asyncio task if active and wait for it to finish
- active = self._active_tasks.get(scenario_result_id)
- if active is not None and active.task is not None and not active.task.done():
- active.task.cancel()
- with contextlib.suppress(asyncio.CancelledError, asyncio.TimeoutError):
- await asyncio.wait_for(active.task, timeout=5.0)
+ task: asyncio.Task[None] | None = None
+ async with self._scheduler_lock:
+ queued = next(
+ (run for run in self._queued_runs if run.scenario_result_id == scenario_result_id),
+ None,
+ )
+ if queued is not None:
+ await asyncio.to_thread(
+ self._memory.update_scenario_run_state,
+ scenario_result_id=scenario_result_id,
+ scenario_run_state=ScenarioRunState.CANCELLED,
+ error_message=_USER_CANCELLATION_REASON,
+ error_type="CancelledError",
+ )
+ self._queued_runs.remove(queued)
+ self._queue_revision += 1
+ elif self._active_scenario_result_id == scenario_result_id:
+ active = self._active_tasks[scenario_result_id]
+ active.cancellation_state = ScenarioRunState.CANCELLED
+ active.cancellation_reason = _USER_CANCELLATION_REASON
+ active.cancellation_error_type = "CancelledError"
+ task = active.task
+ else:
+ latest = await asyncio.to_thread(
+ self._memory.get_scenario_results,
+ scenario_result_ids=[scenario_result_id],
+ )
+ if latest and self._is_terminal_state(latest[0].scenario_run_state):
+ raise ValueError(f"Cannot cancel run in '{latest[0].scenario_run_state}' state.")
+ await asyncio.to_thread(
+ self._memory.update_scenario_run_state,
+ scenario_result_id=scenario_result_id,
+ scenario_run_state=ScenarioRunState.CANCELLED,
+ error_message=_USER_CANCELLATION_REASON,
+ error_type="CancelledError",
+ )
- # Persist cancelled state to DB
- self._memory.update_scenario_run_state(
+ if task is not None and not task.done():
+ task.cancel()
+ with contextlib.suppress(asyncio.CancelledError):
+ await task
+
+ snapshot = self.snapshot_active_run(scenario_result_id=scenario_result_id)
+ result = await asyncio.to_thread(
+ self.get_run_from_storage,
scenario_result_id=scenario_result_id,
- scenario_run_state=ScenarioRunState.CANCELLED,
- error_message="Run was cancelled by user",
- error_type="CancelledError",
+ active_error=snapshot.error,
+ queue_position=snapshot.queue_position,
+ active_scenario_result_id=snapshot.active_scenario_result_id,
+ )
+ if result is not None and result.status != ScenarioRunState.CANCELLED:
+ raise ValueError(f"Cannot cancel run in '{result.status}' state.")
+ return result
+
+ def get_queue_snapshot(self) -> ScenarioQueueSnapshot:
+ """
+ Return the current in-process FIFO scheduler state.
+
+ Returns:
+ ScenarioQueueSnapshot: Active run and ordered waiting runs.
+ """
+ snapshot_at = datetime.now(timezone.utc)
+ active = None
+ if self._active_scenario_result_id is not None:
+ active_run = self._active_tasks.get(self._active_scenario_result_id)
+ if active_run is not None:
+ active = self._build_queue_entry(run=active_run, state=ScenarioRunState.IN_PROGRESS)
+ queued = [
+ self._build_queue_entry(run=run, state=ScenarioRunState.QUEUED, position=position)
+ for position, run in enumerate(self._queued_runs, start=1)
+ ]
+ return ScenarioQueueSnapshot(
+ revision=self._queue_revision,
+ snapshot_at=snapshot_at,
+ active=active,
+ queued=queued,
+ )
+
+ async def reconcile_interrupted_runs_async(self) -> int:
+ """
+ Mark scheduler-managed local rows failed when executable objects were lost.
+
+ Shared and unknown memory backends are intentionally non-destructive because
+ another process may still own their runs. File-backed SQLite assumes one
+ scheduler process has exclusive ownership of that database file.
+
+ Returns:
+ int: Number of reconciled rows.
+ """
+ if not isinstance(self._memory, SQLiteMemory):
+ logger.info(
+ "Skipping interrupted Scenario run reconciliation for shared or unsupported %s memory.",
+ type(self._memory).__name__,
+ )
+ return 0
+
+ states = (ScenarioRunState.CREATED, ScenarioRunState.QUEUED, ScenarioRunState.IN_PROGRESS)
+ after_id = None
+ reconciled = 0
+ while True:
+ interrupted, has_more = await asyncio.to_thread(
+ self._memory.get_scenario_run_state_page,
+ states=states,
+ after_id=after_id,
+ limit=500,
+ )
+ for result in interrupted:
+ header = await asyncio.to_thread(
+ self._memory.get_scenario_result_header,
+ scenario_result_id=result.scenario_result_id,
+ )
+ if header is None or header.metadata.get(_SCHEDULER_METADATA_KEY) != _SCHEDULER_METADATA_VALUE:
+ continue
+ await asyncio.to_thread(
+ self._memory.update_scenario_run_state,
+ scenario_result_id=result.scenario_result_id,
+ scenario_run_state=ScenarioRunState.FAILED,
+ error_message=_RESTART_INTERRUPTION_REASON,
+ error_type=_INTERRUPTED_ERROR_TYPE,
+ )
+ reconciled += 1
+ if not has_more:
+ return reconciled
+ if not interrupted:
+ raise RuntimeError(
+ "Scenario run state projection reported another page without returning a cursor row."
+ )
+ after_id = interrupted[-1].scenario_result_id
+
+ async def shutdown_async(self) -> None:
+ """Stop scheduling and terminalize active and queued runs for process shutdown."""
+ task: asyncio.Task[None] | None = None
+ retry_tasks: list[asyncio.Task[None]] = []
+ errors: list[Exception] = []
+ async with self._scheduler_lock:
+ self._stopping = True
+ retry_tasks = list(self._handoff_retry_tasks)
+ queued = list(self._queued_runs)
+ self._queued_runs.clear()
+ if queued:
+ self._queue_revision += 1
+ for run in queued:
+ try:
+ await asyncio.to_thread(
+ self._memory.update_scenario_run_state,
+ scenario_result_id=run.scenario_result_id,
+ scenario_run_state=ScenarioRunState.FAILED,
+ error_message=_SHUTDOWN_INTERRUPTION_REASON,
+ error_type=_INTERRUPTED_ERROR_TYPE,
+ )
+ except Exception as exc:
+ errors.append(exc)
+ if self._active_scenario_result_id is not None:
+ active = self._active_tasks[self._active_scenario_result_id]
+ active.cancellation_state = ScenarioRunState.FAILED
+ active.cancellation_reason = _SHUTDOWN_INTERRUPTION_REASON
+ active.cancellation_error_type = _INTERRUPTED_ERROR_TYPE
+ task = active.task
+ if task is None or task.done():
+ try:
+ await asyncio.to_thread(
+ self._memory.update_scenario_run_state,
+ scenario_result_id=active.scenario_result_id,
+ scenario_run_state=ScenarioRunState.FAILED,
+ error_message=_SHUTDOWN_INTERRUPTION_REASON,
+ error_type=_INTERRUPTED_ERROR_TYPE,
+ )
+ except Exception as exc:
+ errors.append(exc)
+ self._active_scenario_result_id = None
+ self._release_completed_task(scenario_result_id=active.scenario_result_id)
+ self._queue_revision += 1
+ if task is not None and not task.done():
+ task.cancel()
+ try:
+ await task
+ except asyncio.CancelledError:
+ pass
+ except Exception as exc:
+ errors.append(exc)
+ for retry_task in retry_tasks:
+ retry_task.cancel()
+ if retry_tasks:
+ await asyncio.gather(*retry_tasks, return_exceptions=True)
+ if errors:
+ raise ExceptionGroup("Failed to persist one or more scenario shutdown transitions.", errors)
+
+ async def _enqueue_run_async(self, *, scheduled: _ActiveTask) -> None:
+ """Atomically enqueue a persisted initialized run or start it immediately."""
+ async with self._scheduler_lock:
+ if self._stopping:
+ raise RuntimeError("Scenario run scheduling is stopping.")
+ scheduled_ids = {
+ *(run.scenario_result_id for run in self._queued_runs),
+ *self._active_tasks.keys(),
+ }
+ if scheduled.scenario_result_id in scheduled_ids:
+ raise ValueError(f"Scenario run '{scheduled.scenario_result_id}' is already scheduled.")
+ self._terminal_errors.pop(scheduled.scenario_result_id, None)
+ if self._active_scenario_result_id is None:
+ await self._start_scheduled_run_locked_async(scheduled=scheduled)
+ return
+ await asyncio.to_thread(
+ self._memory.update_scenario_run_state_and_metadata_fields,
+ scenario_result_id=scheduled.scenario_result_id,
+ scenario_run_state=ScenarioRunState.QUEUED,
+ metadata_fields={_SCHEDULER_METADATA_KEY: _SCHEDULER_METADATA_VALUE},
+ )
+ self._queued_runs.append(scheduled)
+ self._queue_revision += 1
+
+ async def _start_scheduled_run_locked_async(self, *, scheduled: _ActiveTask) -> None:
+ """Start one run while the scheduler lock guarantees exclusive ownership."""
+ scheduled.started_at = datetime.now(timezone.utc)
+ await asyncio.to_thread(
+ self._memory.update_scenario_run_state_and_metadata_fields,
+ scenario_result_id=scheduled.scenario_result_id,
+ scenario_run_state=ScenarioRunState.IN_PROGRESS,
+ metadata_fields={
+ _SCHEDULER_METADATA_KEY: _SCHEDULER_METADATA_VALUE,
+ SCENARIO_RUN_STARTED_AT_METADATA_KEY: scheduled.started_at.isoformat(),
+ },
+ )
+ self._active_scenario_result_id = scheduled.scenario_result_id
+ self._active_tasks[scheduled.scenario_result_id] = scheduled
+ scheduled.task = asyncio.create_task(self._execute_run_async(scenario_result_id=scheduled.scenario_result_id))
+ self._queue_revision += 1
+
+ async def _handoff_scheduler_async(self, *, scenario_result_id: str) -> None:
+ """Release one terminal active run and start the next valid queued run once."""
+ async with self._scheduler_lock:
+ if self._active_scenario_result_id != scenario_result_id:
+ return
+ if self._stopping:
+ self._active_scenario_result_id = None
+ self._release_completed_task(scenario_result_id=scenario_result_id)
+ self._queue_revision += 1
+ return
+ while self._queued_runs:
+ next_run = self._queued_runs[0]
+ persisted = await asyncio.to_thread(
+ self._memory.get_scenario_results,
+ scenario_result_ids=[next_run.scenario_result_id],
+ )
+ if not persisted or persisted[0].scenario_run_state != ScenarioRunState.QUEUED:
+ self._queued_runs.popleft()
+ self._queue_revision += 1
+ continue
+ await self._start_scheduled_run_locked_async(scheduled=next_run)
+ self._queued_runs.popleft()
+ self._release_completed_task(scenario_result_id=scenario_result_id)
+ return
+ self._active_scenario_result_id = None
+ self._release_completed_task(scenario_result_id=scenario_result_id)
+ self._queue_revision += 1
+
+ def _release_completed_task(self, *, scenario_result_id: str) -> None:
+ """Release executable state while retaining bounded terminal error evidence."""
+ completed = self._active_tasks.pop(scenario_result_id, None)
+ if completed is None or completed.error is None:
+ return
+ self._terminal_errors[scenario_result_id] = completed.error
+ self._terminal_errors.move_to_end(scenario_result_id)
+ while len(self._terminal_errors) > _MAX_TERMINAL_ERRORS:
+ self._terminal_errors.popitem(last=False)
+
+ def _schedule_handoff_retry(self, *, scenario_result_id: str) -> None:
+ """Retry a failed terminal handoff without permitting another active run."""
+ retry_task = asyncio.create_task(self._retry_handoff_async(scenario_result_id=scenario_result_id))
+ self._handoff_retry_tasks.add(retry_task)
+ retry_task.add_done_callback(self._handoff_retry_tasks.discard)
+
+ def _schedule_terminalization_retry(self, *, active: _ActiveTask) -> None:
+ """Retry cancellation persistence before releasing the active slot."""
+ retry_task = asyncio.create_task(self._retry_terminalization_async(active=active))
+ self._handoff_retry_tasks.add(retry_task)
+ retry_task.add_done_callback(self._handoff_retry_tasks.discard)
+
+ async def _retry_handoff_async(self, *, scenario_result_id: str) -> None:
+ """Retry scheduler handoff with bounded exponential delay until it succeeds or shutdown begins."""
+ delay = _SCHEDULER_RETRY_INITIAL_SECONDS
+ while not self._stopping and self._active_scenario_result_id == scenario_result_id:
+ await asyncio.sleep(delay)
+ try:
+ await self._handoff_scheduler_async(scenario_result_id=scenario_result_id)
+ except Exception:
+ logger.exception("Scenario scheduler handoff retry failed for %s.", scenario_result_id)
+ delay = min(delay * 2, _SCHEDULER_RETRY_MAX_SECONDS)
+ else:
+ return
+
+ async def _retry_terminalization_async(self, *, active: _ActiveTask) -> None:
+ """Retry a failed cancellation transition, then perform the terminal handoff."""
+ delay = _SCHEDULER_RETRY_INITIAL_SECONDS
+ while not self._stopping and self._active_scenario_result_id == active.scenario_result_id:
+ await asyncio.sleep(delay)
+ try:
+ async with self._scheduler_lock:
+ if self._stopping or self._active_scenario_result_id != active.scenario_result_id:
+ return
+ await asyncio.to_thread(
+ self._memory.update_scenario_run_state,
+ scenario_result_id=active.scenario_result_id,
+ scenario_run_state=active.cancellation_state,
+ error_message=active.cancellation_reason,
+ error_type=active.cancellation_error_type,
+ )
+ if not active.retain_error_on_terminalization:
+ active.error = None
+ await self._handoff_scheduler_async(scenario_result_id=active.scenario_result_id)
+ except Exception:
+ logger.exception(
+ "Scenario terminal transition retry failed for %s.",
+ active.scenario_result_id,
+ )
+ delay = min(delay * 2, _SCHEDULER_RETRY_MAX_SECONDS)
+ else:
+ return
+
+ async def _complete_handoff_async(self, *, scenario_result_id: str) -> None:
+ """Complete terminal handoff even if the execution task is cancelled while waiting for the scheduler lock."""
+ handoff_task = asyncio.create_task(self._handoff_scheduler_async(scenario_result_id=scenario_result_id))
+ self._handoff_retry_tasks.add(handoff_task)
+ handoff_task.add_done_callback(self._handoff_retry_tasks.discard)
+ try:
+ await asyncio.shield(handoff_task)
+ except asyncio.CancelledError:
+ try:
+ await handoff_task
+ except asyncio.CancelledError:
+ return
+ except Exception:
+ logger.exception("Scenario scheduler handoff failed for %s; retrying.", scenario_result_id)
+ if not self._stopping:
+ self._schedule_handoff_retry(scenario_result_id=scenario_result_id)
+ except Exception:
+ logger.exception("Scenario scheduler handoff failed for %s; retrying.", scenario_result_id)
+ if not self._stopping:
+ self._schedule_handoff_retry(scenario_result_id=scenario_result_id)
+
+ @staticmethod
+ def _build_queue_entry(
+ *,
+ run: _ActiveTask,
+ state: ScenarioRunState,
+ position: int | None = None,
+ ) -> ScenarioQueueEntry:
+ """
+ Map event-loop scheduler state to the canonical queue DTO.
+
+ Returns:
+ ScenarioQueueEntry: Canonical active or queued entry.
+ """
+ if run.created_at is None or run.enqueued_at is None:
+ raise RuntimeError(f"Scenario run '{run.scenario_result_id}' has incomplete queue timestamps.")
+ return ScenarioQueueEntry(
+ scenario_result_id=run.scenario_result_id,
+ scenario_name=run.scenario_name,
+ scenario_registry_name=run.scenario_registry_name,
+ created_at=run.created_at,
+ enqueued_at=run.enqueued_at,
+ started_at=run.started_at,
+ state=state,
+ position=position,
)
- return self.get_run(scenario_result_id=scenario_result_id)
+ @staticmethod
+ def _is_terminal_state(state: ScenarioRunState) -> bool:
+ """Return whether a scenario state is terminal."""
+ return state in (ScenarioRunState.COMPLETED, ScenarioRunState.FAILED, ScenarioRunState.CANCELLED)
def _resolve_scenario_class(self, *, request: RunScenarioRequest) -> type[Scenario]:
"""
@@ -682,6 +1100,7 @@ async def _initialize_scenario_async(self, *, request: RunScenarioRequest, init_
request.scenario_name,
scenario_params=request.scenario_params or {},
scenario_result_id=request.scenario_result_id or None,
+ initial_metadata={_SCHEDULER_METADATA_KEY: _SCHEDULER_METADATA_VALUE},
**init_kwargs,
)
@@ -691,36 +1110,68 @@ async def _execute_run_async(self, *, scenario_result_id: str) -> None:
Only calls scenario.run_async on the already-initialized scenario.
- Note: this method intentionally does NOT remove the entry from
- ``_active_tasks`` on completion. The entry must stay so that
- ``_build_response_from_db`` can read ``active.error`` when the
- caller next polls the run status. Cleanup happens lazily there
- once the error has been surfaced.
+ Terminal handoff releases executable objects. Bounded error evidence is
+ retained separately for later status polling.
Args:
scenario_result_id: The scenario result ID for this run.
"""
active = self._active_tasks[scenario_result_id]
assert active.scenario is not None
+ handoff_ready = True
try:
await active.scenario.run_async()
except asyncio.CancelledError:
- logger.info(f"Scenario run {scenario_result_id} was cancelled.")
+ try:
+ await asyncio.to_thread(
+ self._memory.update_scenario_run_state,
+ scenario_result_id=scenario_result_id,
+ scenario_run_state=active.cancellation_state,
+ error_message=active.cancellation_reason,
+ error_type=active.cancellation_error_type,
+ )
+ except Exception as exc:
+ handoff_ready = False
+ active.error = str(exc)
+ if not self._stopping:
+ self._schedule_terminalization_retry(active=active)
+ raise
+ logger.info("Scenario run %s stopped in state %s.", scenario_result_id, active.cancellation_state.value)
except Exception as e:
active.error = str(e)
+ active.cancellation_state = ScenarioRunState.FAILED
+ active.cancellation_reason = str(e)
+ active.cancellation_error_type = type(e).__name__
+ active.retain_error_on_terminalization = True
+ try:
+ await asyncio.to_thread(
+ self._memory.update_scenario_run_state,
+ scenario_result_id=scenario_result_id,
+ scenario_run_state=ScenarioRunState.FAILED,
+ error_message=str(e),
+ error_type=type(e).__name__,
+ )
+ except Exception:
+ handoff_ready = False
+ if not self._stopping:
+ self._schedule_terminalization_retry(active=active)
+ logger.exception("Failed to persist terminal state for scenario run %s.", scenario_result_id)
logger.exception(f"Scenario run {scenario_result_id} failed: {e}")
finally:
- self._run_semaphore.release()
+ if handoff_ready:
+ await self._complete_handoff_async(scenario_result_id=scenario_result_id)
def _build_response(
self,
*,
scenario_result_id: str,
active_error: str | None,
+ queue_position: int | None,
+ active_scenario_result_id: str | None,
) -> ScenarioRunSummary | None:
"""
Build a ScenarioRunResponse by querying the database and merging active task state.
@@ -728,6 +1179,8 @@ def _build_response(
Args:
scenario_result_id: The scenario result ID.
active_error: Error copied from the active asyncio task, if any.
+ queue_position: Current 1-based waiting position, if queued.
+ active_scenario_result_id: Currently executing scenario result ID.
Returns:
ScenarioRunResponse if found in the database, None otherwise.
@@ -735,13 +1188,20 @@ def _build_response(
results = self._memory.get_scenario_results(scenario_result_ids=[scenario_result_id])
if not results:
return None
- return self._build_response_from_db(scenario_result=results[0], active_error=active_error)
+ return self._build_response_from_db(
+ scenario_result=results[0],
+ active_error=active_error,
+ queue_position=queue_position,
+ active_scenario_result_id=active_scenario_result_id,
+ )
def _build_response_from_db(
self,
*,
scenario_result: ScenarioResult,
active_error: str | None = None,
+ queue_position: int | None = None,
+ active_scenario_result_id: str | None = None,
) -> ScenarioRunSummary:
"""
Build a ScenarioRunResponse from a database ScenarioResult, merged with active task info.
@@ -749,6 +1209,8 @@ def _build_response_from_db(
Args:
scenario_result: A ScenarioResult retrieved from CentralMemory.
active_error: Error copied from the active asyncio task, if any.
+ queue_position: Current 1-based waiting position, if queued.
+ active_scenario_result_id: Currently executing scenario result ID.
Returns:
The API response model.
@@ -807,7 +1269,8 @@ def _build_response_from_db(
# a COMPLETED scenario can still hide errored objectives or rate-limit retries.
failed_attacks: list[AttackErrorSummary] = []
attack_retries: list[AttackRetrySummary] = []
- total_retries = 0
+ overload_events: deque[Any] = deque(maxlen=_MAX_OVERLOAD_EVENTS)
+ inner_retries_by_unit: dict[tuple[str, str], int] = {}
attempts_by_unit: dict[tuple[str, str], int] = {}
for atomic_attack_name, results in scenario_result.attack_results.items():
for attack_result in results:
@@ -819,10 +1282,11 @@ def _build_response_from_db(
attempts_by_unit[unit_key] = attempts_by_unit.get(unit_key, 0) + 1
retries = getattr(attack_result, "total_retries", 0)
if isinstance(retries, int):
- total_retries += retries
+ inner_retries_by_unit[unit_key] = inner_retries_by_unit.get(unit_key, 0) + max(0, retries)
retry_events = getattr(attack_result, "retry_events", None)
if isinstance(retry_events, list) and retry_events:
+ overload_events.extend(retry_events)
attack_retries.append(
AttackRetrySummary(
attack_result_id=str(attack_result.attack_result_id),
@@ -838,10 +1302,16 @@ def _build_response_from_db(
objective=attack_result.objective,
error_type=attack_result.error_type,
error_message=attack_result.error_message,
- total_retries=retries if isinstance(retries, int) else 0,
+ total_retries=max(0, retries) if isinstance(retries, int) else 0,
)
)
- total_retries += sum(max(0, attempt_count - 1) for attempt_count in attempts_by_unit.values())
+ total_retries = sum(
+ self._total_retry_work(
+ inner_retries=inner_retries_by_unit.get(unit_key, 0),
+ attempt_count=attempt_count,
+ )
+ for unit_key, attempt_count in attempts_by_unit.items()
+ )
updated_at = scenario_result.creation_time
if terminal and scenario_result.completion_time is not None:
@@ -854,6 +1324,7 @@ def _build_response_from_db(
scenario_version=scenario_result.scenario_version,
status=status,
created_at=scenario_result.creation_time,
+ started_at=self._load_started_at(scenario_result=scenario_result),
updated_at=updated_at,
error=error,
error_type=error_type,
@@ -877,6 +1348,9 @@ def _build_response_from_db(
planned_total_available=plan is not None,
successful_attacks=successful_attacks,
error_attacks=len(failed_attacks),
+ queue_position=queue_position,
+ active_scenario_result_id=active_scenario_result_id,
+ overload_summaries=self._build_overload_summaries(retry_events=overload_events),
)
def _build_history_summary(
@@ -993,6 +1467,7 @@ def _build_history_summary(
scenario_version=record.scenario_version,
status=status,
created_at=record.created_at,
+ started_at=record.started_at,
updated_at=updated_at,
error=record.error_message,
error_type=record.error_type,
@@ -1056,6 +1531,9 @@ def _merge_history_units(
preferred = incoming if incoming_completed else existing
else:
preferred = incoming if incoming.latest_timestamp > existing.latest_timestamp else existing
+ attempt_count = max(0, existing.attempt_count) + max(0, incoming.attempt_count)
+ existing_inner_retries = max(0, existing.total_retries - max(0, existing.attempt_count - 1))
+ incoming_inner_retries = max(0, incoming.total_retries - max(0, incoming.attempt_count - 1))
return ScenarioHistoryUnitRecord(
scenario_result_id=preferred.scenario_result_id,
atomic_attack_name=preferred.atomic_attack_name,
@@ -1064,10 +1542,41 @@ def _merge_history_units(
objective_sha256=preferred.objective_sha256 or existing.objective_sha256 or incoming.objective_sha256,
latest_outcome=preferred.latest_outcome,
latest_timestamp=max(existing.latest_timestamp, incoming.latest_timestamp),
- total_retries=max(0, existing.total_retries) + max(0, incoming.total_retries) + 1,
+ total_retries=ScenarioRunService._total_retry_work(
+ inner_retries=existing_inner_retries + incoming_inner_retries,
+ attempt_count=attempt_count,
+ ),
error_count=max(0, existing.error_count) + max(0, incoming.error_count),
+ attempt_count=attempt_count,
)
+ @staticmethod
+ def _total_retry_work(*, inner_retries: int, attempt_count: int) -> int:
+ """
+ Count retry work beyond the first logical attempt.
+
+ Returns:
+ int: Inner retries plus additional scenario attempts.
+ """
+ return max(0, inner_retries) + max(0, attempt_count - 1)
+
+ @staticmethod
+ def _load_started_at(*, scenario_result: ScenarioResult) -> datetime | None:
+ """
+ Load the persisted aware execution start timestamp from scenario metadata.
+
+ Returns:
+ datetime | None: The execution start, or None for legacy or invalid metadata.
+ """
+ raw_value = (getattr(scenario_result, "metadata", None) or {}).get(SCENARIO_RUN_STARTED_AT_METADATA_KEY)
+ if raw_value is None:
+ return None
+ try:
+ started_at = _STARTED_AT_ADAPTER.validate_python(raw_value)
+ except ValidationError:
+ return None
+ return started_at if started_at.tzinfo is not None else None
+
@staticmethod
def _safe_run_metadata(
*,
@@ -1090,6 +1599,55 @@ def _safe_run_metadata(
ScenarioRunService._safe_scenario_parameters(parameters=dict(scenario_identifier.params)),
)
+ @staticmethod
+ def _build_overload_summaries(*, retry_events: Sequence[Any]) -> list[ScenarioOverloadSummary]:
+ """
+ Aggregate bounded HTTP overload evidence by component role.
+
+ Returns:
+ list[ScenarioOverloadSummary]: Most recently affected roles first.
+ """
+ aggregates: dict[str, dict[str, Any]] = {}
+ for event in retry_events:
+ status_code = getattr(event, "status_code", None)
+ if not isinstance(status_code, int) or (status_code != 429 and not 500 <= status_code <= 599):
+ continue
+ role = str(getattr(event, "component_role", "") or "unknown")
+ timestamp = getattr(event, "timestamp", None)
+ if not isinstance(timestamp, datetime):
+ continue
+ aggregate = aggregates.setdefault(
+ role,
+ {
+ "count": 0,
+ "rate_limit_count": 0,
+ "server_error_count": 0,
+ "status_codes": set(),
+ "latest_timestamp": timestamp,
+ },
+ )
+ aggregate["count"] += 1
+ aggregate["rate_limit_count"] += status_code == 429
+ aggregate["server_error_count"] += 500 <= status_code <= 599
+ aggregate["status_codes"].add(status_code)
+ aggregate["latest_timestamp"] = max(aggregate["latest_timestamp"], timestamp)
+ ordered = sorted(
+ aggregates.items(),
+ key=lambda item: item[1]["latest_timestamp"],
+ reverse=True,
+ )[:_MAX_OVERLOAD_ROLES]
+ return [
+ ScenarioOverloadSummary(
+ component_role=role,
+ count=aggregate["count"],
+ rate_limit_count=aggregate["rate_limit_count"],
+ server_error_count=aggregate["server_error_count"],
+ status_codes=sorted(aggregate["status_codes"]),
+ latest_timestamp=aggregate["latest_timestamp"],
+ )
+ for role, aggregate in ordered
+ ]
+
@staticmethod
def _safe_target_metadata(*, target_identifier: TargetIdentifier | None) -> ScenarioTargetSummary | None:
"""
@@ -1251,10 +1809,16 @@ def _decode_history_cursor(
)
def _get_active_task(self, *, scenario_result_id: str) -> _ActiveTask | None:
- """Return a live task and release completed task state."""
+ """Return executable state for an active run."""
active = self._active_tasks.get(scenario_result_id)
- if active is not None and active.task is not None and active.task.done():
- self._active_tasks.pop(scenario_result_id, None)
+ if (
+ active is not None
+ and active.task is not None
+ and active.task.done()
+ and self._active_scenario_result_id != scenario_result_id
+ ):
+ self._release_completed_task(scenario_result_id=scenario_result_id)
+ return None
return active
def snapshot_active_run(self, *, scenario_result_id: str) -> _ActiveRunSnapshot:
@@ -1264,11 +1828,29 @@ def snapshot_active_run(self, *, scenario_result_id: str) -> _ActiveRunSnapshot:
Returns:
_ActiveRunSnapshot: An immutable copy of the active state.
"""
+ active_scenario_result_id = self._active_scenario_result_id
+ queue_position = next(
+ (
+ position
+ for position, queued in enumerate(self._queued_runs, start=1)
+ if queued.scenario_result_id == scenario_result_id
+ ),
+ None,
+ )
active = self._get_active_task(scenario_result_id=scenario_result_id)
if active is None:
- return _ActiveRunSnapshot()
+ return _ActiveRunSnapshot(
+ error=self._terminal_errors.get(scenario_result_id),
+ queue_position=queue_position,
+ active_scenario_result_id=active_scenario_result_id,
+ )
active_group_ids = tuple(sorted(active.scenario.active_atomic_group_ids)) if active.scenario is not None else ()
- return _ActiveRunSnapshot(error=active.error, active_group_ids=active_group_ids)
+ return _ActiveRunSnapshot(
+ error=active.error,
+ active_group_ids=active_group_ids,
+ queue_position=queue_position,
+ active_scenario_result_id=active_scenario_result_id,
+ )
@staticmethod
def _load_run_plan(*, scenario_result: ScenarioResult) -> ScenarioRunPlan | None:
@@ -1401,6 +1983,8 @@ def get_run_progress(
since=since,
limit=limit,
active_group_ids=snapshot.active_group_ids,
+ queue_position=snapshot.queue_position,
+ active_scenario_result_id=snapshot.active_scenario_result_id,
)
def get_run_progress_from_storage(
@@ -1410,6 +1994,8 @@ def get_run_progress_from_storage(
since: str | None,
limit: int,
active_group_ids: Sequence[str],
+ queue_position: int | None = None,
+ active_scenario_result_id: str | None = None,
) -> ScenarioRunProgress | None:
"""Return compact database progress using a previously captured live-state snapshot."""
header_result = self._memory.get_scenario_result_header(scenario_result_id=scenario_result_id)
@@ -1429,6 +2015,9 @@ def get_run_progress_from_storage(
response_plan = self._synthesize_legacy_plan(deltas=deltas)
results = [self._map_progress_delta(delta=delta, plan=plan or response_plan) for delta in deltas]
+ overload_events: deque[Any] = deque(maxlen=_MAX_OVERLOAD_EVENTS)
+ for delta in deltas:
+ overload_events.extend(delta.retry_events)
next_cursor = (
self._encode_progress_cursor(scenario_result_id=scenario_result_id, delta=deltas[-1]) if deltas else since
)
@@ -1453,6 +2042,7 @@ def get_run_progress_from_storage(
scenario_version=header_result.scenario_version,
status=header_result.scenario_run_state,
created_at=header_result.creation_time,
+ started_at=self._load_started_at(scenario_result=header_result),
completed_at=header_result.completion_time if terminal else None,
pyrit_version=header_result.pyrit_version,
target=target,
@@ -1460,6 +2050,9 @@ def get_run_progress_from_storage(
datasets_used=datasets_used,
scenario_parameters=scenario_parameters,
labels=header_result.labels,
+ queue_position=queue_position,
+ active_scenario_result_id=active_scenario_result_id,
+ overload_summaries=self._build_overload_summaries(retry_events=overload_events),
),
plan=response_plan,
reset=False,
diff --git a/pyrit/exceptions/retry_collector.py b/pyrit/exceptions/retry_collector.py
index 4f5d216498..703b9e085f 100644
--- a/pyrit/exceptions/retry_collector.py
+++ b/pyrit/exceptions/retry_collector.py
@@ -41,12 +41,18 @@ def record(self, *, retry_state: RetryCallState) -> None:
# Extract exception info
exception_type = ""
exception_message = ""
+ status_code: int | None = None
outcome = retry_state.outcome
if outcome is not None and outcome.failed:
exc = outcome.exception()
if exc:
exception_type = type(exc).__name__
exception_message = str(exc)
+ candidate_status = getattr(exc, "status_code", None)
+ if not isinstance(candidate_status, int):
+ candidate_status = getattr(getattr(exc, "response", None), "status_code", None)
+ if isinstance(candidate_status, int):
+ status_code = candidate_status
# Extract context info
component_role = ""
@@ -69,6 +75,7 @@ def record(self, *, retry_state: RetryCallState) -> None:
component_role=component_role,
component_name=component_name,
endpoint=endpoint,
+ status_code=status_code,
elapsed_seconds=round(elapsed, 3),
)
self.events.append(event)
diff --git a/pyrit/memory/__init__.py b/pyrit/memory/__init__.py
index efd2d82301..080203d370 100644
--- a/pyrit/memory/__init__.py
+++ b/pyrit/memory/__init__.py
@@ -17,6 +17,7 @@
ScenarioHistoryRunRecord,
ScenarioHistoryUnitRecord,
ScenarioProgressKeysetCursor,
+ ScenarioRunStateRecord,
)
from pyrit.memory.memory_models import AttackResultEntry, EmbeddingDataEntry, PromptMemoryEntry, SeedEntry
from pyrit.memory.sqlite_memory import SQLiteMemory
@@ -60,6 +61,7 @@
"ScenarioHistoryRunRecord",
"ScenarioHistoryUnitRecord",
"ScenarioProgressKeysetCursor",
+ "ScenarioRunStateRecord",
"PromptMemoryEntry",
"SeedEntry",
"set_message_piece_sha256_async",
diff --git a/pyrit/memory/azure_sql_memory.py b/pyrit/memory/azure_sql_memory.py
index 272db3093f..6da97e7a86 100644
--- a/pyrit/memory/azure_sql_memory.py
+++ b/pyrit/memory/azure_sql_memory.py
@@ -590,7 +590,25 @@ def get_conversation_stats(self, *, conversation_ids: Sequence[str]) -> dict[str
return result
- def _get_scenario_result_label_condition(self, *, labels: Mapping[str, str | Sequence[str]]) -> Any:
+ def _get_scenario_result_label_condition(self, *, labels: dict[str, str]) -> Any:
+ """
+ Filter ScenarioResults by legacy single-value labels.
+
+ Returns:
+ Any: SQLAlchemy condition for all supplied labels.
+ """
+ conditions = []
+ for key_index, (key, value) in enumerate(labels.items()):
+ path_param = f"scenario_label_path_{key_index}"
+ value_param = f"scenario_label_value_{key_index}"
+ conditions.append(
+ text(f"ISJSON(labels) = 1 AND JSON_VALUE(labels, :{path_param}) = :{value_param}").bindparams(
+ **{path_param: f'$."{key}"', value_param: value}
+ )
+ )
+ return and_(*conditions)
+
+ def _get_scenario_result_labels_condition(self, *, labels: Mapping[str, str | Sequence[str]]) -> Any:
"""
Get the SQL Azure implementation for filtering ScenarioResults by labels.
@@ -667,6 +685,10 @@ def _get_scenario_history_plan_expressions(self) -> tuple[Any, Any, Any]:
),
)
+ def _get_scenario_started_at_expression(self) -> Any:
+ """Return the persisted execution start without loading full scenario metadata."""
+ return func.json_value(ScenarioResultEntry.scenario_metadata, "$.started_at")
+
def _get_scenario_attempt_unit_expressions(self) -> tuple[Any, Any, Any]:
"""Return SQL Server JSON expressions for persisted scenario attempt attribution."""
atomic_name = func.coalesce(
diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py
index c7c26a51be..3137bf61d1 100644
--- a/pyrit/memory/memory_interface.py
+++ b/pyrit/memory/memory_interface.py
@@ -15,7 +15,7 @@
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, ClassVar, Literal, NamedTuple, TypeVar
-from sqlalchemy import MetaData, and_, case, func, not_, or_, select
+from sqlalchemy import MetaData, and_, case, func, literal, not_, or_, select
from sqlalchemy.engine.base import Engine
from sqlalchemy.exc import IntegrityError, SQLAlchemyError
from sqlalchemy.orm import joinedload
@@ -158,11 +158,17 @@ class ScenarioHistoryRunRecord:
scenario_registry_name: str | None
plan_atomic_groups: str | list[dict[str, Any]] | None
plan_seed_id_map: str | list[dict[str, str]] | None
+ started_at: datetime | None = None
@dataclass(frozen=True, slots=True, kw_only=True)
class ScenarioHistoryUnitRecord:
- """One logical scenario work unit aggregated from all persisted attempts."""
+ """
+ One logical scenario work unit aggregated from all persisted attempts.
+
+ ``total_retries`` is all work beyond the initial logical attempt: inner
+ retries plus additional scenario-level attempts.
+ """
scenario_result_id: str
atomic_attack_name: str
@@ -173,6 +179,15 @@ class ScenarioHistoryUnitRecord:
latest_timestamp: datetime
total_retries: int
error_count: int
+ attempt_count: int = 1
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class ScenarioRunStateRecord:
+ """Lightweight persisted ID/state projection used for startup reconciliation."""
+
+ scenario_result_id: str
+ state: ScenarioRunState
@dataclass(frozen=True, slots=True, kw_only=True)
@@ -1637,17 +1652,38 @@ def get_conversation_stats(self, *, conversation_ids: Sequence[str]) -> dict[str
"""
@abc.abstractmethod
- def _get_scenario_result_label_condition(self, *, labels: Mapping[str, str | Sequence[str]]) -> Any:
+ def _get_scenario_result_label_condition(self, *, labels: dict[str, str]) -> Any:
"""
Return a database-specific condition for filtering ScenarioResults by labels.
Args:
- labels: Labels with OR-within-key and AND-across-key semantics.
+ labels: Legacy single-value labels with AND-across-key semantics.
Returns:
Database-specific SQLAlchemy condition.
"""
+ def _get_scenario_result_labels_condition(self, *, labels: Mapping[str, str | Sequence[str]]) -> Any:
+ """
+ Compose multi-value label filters through the legacy single-value hook.
+
+ Returns:
+ Any: OR-within-key and AND-across-key SQLAlchemy condition.
+ """
+ conditions = []
+ for key, raw_value in labels.items():
+ values = [raw_value] if isinstance(raw_value, str) else list(raw_value)
+ if values:
+ conditions.append(
+ or_(
+ *(
+ self._get_scenario_result_label_condition(labels={key: str(value)}).unique_params()
+ for value in values
+ )
+ )
+ )
+ return and_(*conditions)
+
def _get_scenario_registry_name_condition(self, *, scenario_names: Sequence[str]) -> Any:
"""
Return a backend-specific condition matching persisted run-plan registry names.
@@ -1672,6 +1708,10 @@ def _get_scenario_history_plan_expressions(self) -> tuple[Any, Any, Any]:
"to support Scenario history queries."
)
+ def _get_scenario_started_at_expression(self) -> Any:
+ """Return a compact persisted start-time expression when the backend supports one."""
+ return literal(None)
+
def _get_scenario_attempt_unit_expressions(self) -> tuple[Any, Any, Any]:
"""
Return backend-specific JSON expressions for scenario attempt unit attribution.
@@ -3396,6 +3436,29 @@ def update_scenario_run_state(
error_message (str | None): Optional scenario-level error message.
error_type (str | None): Optional exception class name.
+ Raises:
+ ValueError: If the scenario result is not found.
+ """
+ self.update_scenario_run_state_and_metadata_fields(
+ scenario_result_id=scenario_result_id,
+ scenario_run_state=scenario_run_state,
+ error_message=error_message,
+ error_type=error_type,
+ metadata_fields={},
+ )
+
+ def update_scenario_run_state_and_metadata_fields(
+ self,
+ *,
+ scenario_result_id: str,
+ scenario_run_state: ScenarioRunState,
+ metadata_fields: Mapping[str, Any],
+ error_message: str | None = None,
+ error_type: str | None = None,
+ ) -> None:
+ """
+ Update run state and merge scenario metadata in one transaction.
+
Raises:
ValueError: If the scenario result is not found.
"""
@@ -3408,6 +3471,9 @@ def update_scenario_run_state(
entry.scenario_run_state = scenario_run_state.value
entry.error_message = error_message
entry.error_type = error_type
+ if metadata_fields:
+ entry.scenario_metadata = {**(entry.scenario_metadata or {}), **metadata_fields}
+ flag_modified(entry, "scenario_metadata")
if scenario_run_state in (
ScenarioRunState.COMPLETED,
ScenarioRunState.FAILED,
@@ -3447,12 +3513,72 @@ def update_scenario_metadata(
entry.scenario_metadata = metadata if metadata else None
session.commit()
+ def update_scenario_metadata_fields(
+ self,
+ *,
+ scenario_result_id: str,
+ fields: Mapping[str, Any],
+ ) -> None:
+ """
+ Merge selected fields into persisted scenario metadata in one transaction.
+
+ Raises:
+ ValueError: If the scenario result is not found.
+ """
+ with closing(self.get_session()) as session:
+ entry = session.query(ScenarioResultEntry).filter_by(id=scenario_result_id).first()
+ if not entry:
+ raise ValueError(f"Scenario result with ID {scenario_result_id} not found in memory")
+ entry.scenario_metadata = {**(entry.scenario_metadata or {}), **fields}
+ flag_modified(entry, "scenario_metadata")
+ session.commit()
+
def get_scenario_result_header(self, *, scenario_result_id: str) -> ScenarioResult | None:
"""Return one ScenarioResult header without hydrating linked attack results."""
with closing(self.get_session()) as session:
entry = session.query(ScenarioResultEntry).filter_by(id=scenario_result_id).first()
return entry.get_scenario_result() if entry is not None else None
+ def get_scenario_run_state_page(
+ self,
+ *,
+ states: Sequence[ScenarioRunState],
+ after_id: str | None = None,
+ limit: int = 500,
+ ) -> tuple[list[ScenarioRunStateRecord], bool]:
+ """
+ Return one bounded ID/state page without hydrating ScenarioResults or AttackResults.
+
+ Returns:
+ tuple[list[ScenarioRunStateRecord], bool]: State records and whether another page exists.
+
+ Raises:
+ ValueError: If the limit or cursor ID is invalid.
+ """
+ if limit < 1 or limit > 500:
+ raise ValueError("Scenario run state projection limit must be between 1 and 500.")
+ conditions = [ScenarioResultEntry.scenario_run_state.in_([state.value for state in states])]
+ if after_id is not None:
+ conditions.append(ScenarioResultEntry.id > uuid.UUID(after_id))
+ statement = (
+ select(ScenarioResultEntry.id, ScenarioResultEntry.scenario_run_state)
+ .where(and_(*conditions))
+ .order_by(ScenarioResultEntry.id.asc())
+ .limit(limit + 1)
+ )
+ with closing(self.get_session()) as session:
+ rows = session.execute(statement).all()
+ return (
+ [
+ ScenarioRunStateRecord(
+ scenario_result_id=str(row.id),
+ state=ScenarioRunState(row.scenario_run_state),
+ )
+ for row in rows[:limit]
+ ],
+ len(rows) > limit,
+ )
+
def get_scenario_run_history_page(
self,
*,
@@ -3503,7 +3629,7 @@ def get_scenario_run_history_page(
f"Invalid label key(s) {invalid_keys!r}: keys must match {self._LABEL_KEY_PATTERN.pattern}."
)
if effective_labels:
- conditions.append(self._get_scenario_result_label_condition(labels=effective_labels))
+ conditions.append(self._get_scenario_result_labels_condition(labels=effective_labels))
if cursor is not None:
cursor_id = uuid.UUID(cursor.scenario_result_id)
conditions.append(
@@ -3526,6 +3652,7 @@ def get_scenario_run_history_page(
ScenarioResultEntry.scenario_run_state,
ScenarioResultEntry.labels,
ScenarioResultEntry.timestamp,
+ self._get_scenario_started_at_expression().label("started_at"),
ScenarioResultEntry.completion_time,
ScenarioResultEntry.error_message,
ScenarioResultEntry.error_type,
@@ -3565,11 +3692,18 @@ def get_scenario_run_history_page(
AttackResultEntry.objective_sha256.label("objective_sha256"),
AttackResultEntry.outcome.label("latest_outcome"),
func.max(AttackResultEntry.timestamp).over(partition_by=unit_partition).label("latest_timestamp"),
- (
- func.sum(func.coalesce(AttackResultEntry.total_retries, 0)).over(partition_by=unit_partition)
- + func.count().over(partition_by=unit_partition)
- - 1
- ).label("total_retries"),
+ func.sum(
+ case(
+ (
+ func.coalesce(AttackResultEntry.total_retries, 0) > 0,
+ func.coalesce(AttackResultEntry.total_retries, 0),
+ ),
+ else_=0,
+ )
+ )
+ .over(partition_by=unit_partition)
+ .label("inner_retries"),
+ func.count().over(partition_by=unit_partition).label("attempt_count"),
func.sum(case((AttackResultEntry.outcome == AttackOutcome.ERROR.value, 1), else_=0))
.over(partition_by=unit_partition)
.label("error_count"),
@@ -3598,6 +3732,7 @@ def get_scenario_run_history_page(
status=row.scenario_run_state,
labels=row.labels or {},
created_at=row.timestamp,
+ started_at=self._parse_scenario_started_at(raw_value=row.started_at),
completed_at=row.completion_time,
error_message=row.error_message,
error_type=row.error_type,
@@ -3621,12 +3756,29 @@ def get_scenario_run_history_page(
objective_sha256=row.objective_sha256,
latest_outcome=row.latest_outcome,
latest_timestamp=row.latest_timestamp,
- total_retries=row.total_retries or 0,
+ total_retries=(row.inner_retries or 0) + max(0, (row.attempt_count or 0) - 1),
error_count=row.error_count or 0,
+ attempt_count=row.attempt_count or 0,
)
)
return records, units_by_run, len(rows) > limit
+ @staticmethod
+ def _parse_scenario_started_at(*, raw_value: Any) -> datetime | None:
+ """
+ Parse a persisted aware scenario start timestamp.
+
+ Returns:
+ datetime | None: Aware start timestamp, or None for legacy or malformed values.
+ """
+ if not isinstance(raw_value, str):
+ return None
+ try:
+ value = datetime.fromisoformat(raw_value)
+ except ValueError:
+ return None
+ return value if value.tzinfo is not None else None
+
def get_unique_scenario_labels(self) -> dict[str, list[str]]:
"""Return all unique label values across scenario results."""
label_values: dict[str, set[str]] = {}
diff --git a/pyrit/memory/sqlite_memory.py b/pyrit/memory/sqlite_memory.py
index 74f41bc60a..3b2f5651e9 100644
--- a/pyrit/memory/sqlite_memory.py
+++ b/pyrit/memory/sqlite_memory.py
@@ -439,9 +439,20 @@ def get_conversation_stats(self, *, conversation_ids: Sequence[str]) -> dict[str
return result
- def _get_scenario_result_label_condition(self, *, labels: Mapping[str, str | Sequence[str]]) -> Any:
+ def _get_scenario_result_label_condition(self, *, labels: dict[str, str]) -> Any:
"""
- SQLite implementation for filtering ScenarioResults by labels.
+ Filter ScenarioResults by legacy single-value labels.
+
+ Returns:
+ Any: SQLAlchemy condition for all supplied labels.
+ """
+ return and_(
+ *(func.json_extract(ScenarioResultEntry.labels, f'$."{key}"') == value for key, value in labels.items())
+ )
+
+ def _get_scenario_result_labels_condition(self, *, labels: Mapping[str, str | Sequence[str]]) -> Any:
+ """
+ SQLite implementation for filtering ScenarioResults by multi-value labels.
Uses json_extract() function specific to SQLite.
Returns:
@@ -499,6 +510,10 @@ def _get_scenario_history_plan_expressions(self) -> tuple[Any, Any, Any]:
compact_seed_map,
)
+ def _get_scenario_started_at_expression(self) -> Any:
+ """Return the persisted execution start without loading full scenario metadata."""
+ return func.json_extract(ScenarioResultEntry.scenario_metadata, "$.started_at")
+
def _get_scenario_attempt_unit_expressions(self) -> tuple[Any, Any, Any]:
"""Return SQLite JSON expressions for persisted scenario attempt attribution."""
atomic_name = func.coalesce(
diff --git a/pyrit/models/__init__.py b/pyrit/models/__init__.py
index c73ce8f804..4441401427 100644
--- a/pyrit/models/__init__.py
+++ b/pyrit/models/__init__.py
@@ -103,9 +103,12 @@
from pyrit.models.scenario_progress import (
SCENARIO_RUN_PLAN_METADATA_KEY,
SCENARIO_RUN_PLAN_VERSION,
+ SCENARIO_RUN_STARTED_AT_METADATA_KEY,
ScenarioAttackResultDelta,
ScenarioProgressHeader,
ScenarioProgressResult,
+ ScenarioQueueEntry,
+ ScenarioQueueSnapshot,
ScenarioRunPlan,
ScenarioRunPlanAtomicGroup,
ScenarioRunPlanGroupKind,
@@ -235,8 +238,11 @@
"ScenarioRunSizeFactor",
"ScenarioRunState",
"SCENARIO_RUN_PLAN_METADATA_KEY",
+ "SCENARIO_RUN_STARTED_AT_METADATA_KEY",
"SCENARIO_RUN_PLAN_VERSION",
"ScenarioAttackResultDelta",
+ "ScenarioQueueEntry",
+ "ScenarioQueueSnapshot",
"ScenarioProgressHeader",
"ScenarioProgressResult",
"ScenarioRunPlan",
diff --git a/pyrit/models/catalog/scenario.py b/pyrit/models/catalog/scenario.py
index 4c78d477d9..ee19a54ee5 100644
--- a/pyrit/models/catalog/scenario.py
+++ b/pyrit/models/catalog/scenario.py
@@ -374,6 +374,17 @@ class AttackRetrySummary(BaseModel):
)
+class ScenarioOverloadSummary(BaseModel):
+ """Recent structured overload signals grouped by component role."""
+
+ component_role: str = Field(..., description="Role of the component that observed overload")
+ count: int = Field(..., ge=1, description="Recent HTTP 429 and 5xx retry signals")
+ rate_limit_count: int = Field(0, ge=0, description="Recent HTTP 429 retry signals")
+ server_error_count: int = Field(0, ge=0, description="Recent HTTP 5xx retry signals")
+ status_codes: list[int] = Field(default_factory=list, description="Observed overload status codes")
+ latest_timestamp: datetime = Field(..., description="Latest overload signal timestamp")
+
+
class ScenarioRunSummary(BaseModel):
"""Response for a scenario run (status + result details)."""
@@ -383,6 +394,7 @@ class ScenarioRunSummary(BaseModel):
scenario_version: int = Field(0, ge=0, description="Version of the scenario")
status: ScenarioRunState = Field(..., description="Current run status")
created_at: datetime = Field(..., description="When the run was created")
+ started_at: datetime | None = Field(None, description="When active scenario execution started")
updated_at: datetime = Field(..., description="When the run status last changed")
error: str | None = Field(None, description="Error message if status is FAILED")
error_type: str | None = Field(None, description="Exception class name if status is FAILED")
@@ -399,7 +411,10 @@ class ScenarioRunSummary(BaseModel):
description="Per-attack retry events, surfaced as each attack result lands so the CLI can stream warnings",
)
total_retries: int = Field(
- 0, ge=0, description="Total retry attempts recorded across all attack results (endpoint-stress signal)"
+ 0,
+ ge=0,
+ description="Total retry work beyond each logical unit's initial attempt, including inner retries "
+ "and additional scenario attempts",
)
labels: dict[str, str] = Field(default_factory=dict, description="Labels attached to this run")
completed_at: datetime | None = Field(None, description="When the scenario finished")
@@ -420,6 +435,12 @@ class ScenarioRunSummary(BaseModel):
True,
description="Whether failed_attacks and attack_retries contain per-attempt details",
)
+ queue_position: int | None = Field(None, ge=1, description="Current 1-based waiting position")
+ active_scenario_result_id: str | None = Field(None, description="Currently executing scenario result ID")
+ overload_summaries: list[ScenarioOverloadSummary] = Field(
+ default_factory=list,
+ description="Bounded recent HTTP 429 and 5xx retry evidence grouped by component role",
+ )
class ScenarioTargetSummary(BaseModel):
diff --git a/pyrit/models/results/scenario_result.py b/pyrit/models/results/scenario_result.py
index bddfb26f20..bb03ec8a8b 100644
--- a/pyrit/models/results/scenario_result.py
+++ b/pyrit/models/results/scenario_result.py
@@ -46,6 +46,7 @@ class ScenarioRunState(str, Enum):
"""
CREATED = "CREATED"
+ QUEUED = "QUEUED"
IN_PROGRESS = "IN_PROGRESS"
COMPLETED = "COMPLETED"
FAILED = "FAILED"
diff --git a/pyrit/models/retry_event.py b/pyrit/models/retry_event.py
index 7f5a0f7982..6e0945372d 100644
--- a/pyrit/models/retry_event.py
+++ b/pyrit/models/retry_event.py
@@ -28,4 +28,5 @@ class RetryEvent(BaseModel):
component_role: str = ""
component_name: str | None = None
endpoint: str | None = None
+ status_code: int | None = None
elapsed_seconds: float = 0.0
diff --git a/pyrit/models/scenario_progress.py b/pyrit/models/scenario_progress.py
index 341f104135..4ef9405057 100644
--- a/pyrit/models/scenario_progress.py
+++ b/pyrit/models/scenario_progress.py
@@ -9,13 +9,14 @@
from pydantic import AwareDatetime, BaseModel, Field, model_validator
-from pyrit.models.catalog.scenario import ScenarioTargetSummary # noqa: TC001
+from pyrit.models.catalog.scenario import ScenarioOverloadSummary, ScenarioTargetSummary # noqa: TC001
from pyrit.models.identifiers.atomic_attack_identifier import AtomicAttackIdentifier
from pyrit.models.results.attack_result import AttackOutcome
from pyrit.models.results.scenario_result import ScenarioRunState
from pyrit.models.retry_event import RetryEvent
SCENARIO_RUN_PLAN_METADATA_KEY = "run_plan"
+SCENARIO_RUN_STARTED_AT_METADATA_KEY = "started_at"
SCENARIO_RUN_PLAN_VERSION = 1
@@ -95,6 +96,7 @@ class ScenarioProgressHeader(BaseModel):
scenario_version: int
status: ScenarioRunState
created_at: datetime
+ started_at: AwareDatetime | None = None
completed_at: datetime | None = None
pyrit_version: str | None = None
target: "ScenarioTargetSummary | None" = None
@@ -102,6 +104,9 @@ class ScenarioProgressHeader(BaseModel):
datasets_used: list[str] = Field(default_factory=list)
scenario_parameters: dict[str, Any] = Field(default_factory=dict)
labels: dict[str, str] = Field(default_factory=dict)
+ queue_position: int | None = Field(None, ge=1)
+ active_scenario_result_id: str | None = None
+ overload_summaries: list["ScenarioOverloadSummary"] = Field(default_factory=list)
class ScenarioProgressResult(BaseModel):
@@ -133,6 +138,28 @@ class ScenarioRunProgress(BaseModel):
plan_complete: bool
+class ScenarioQueueEntry(BaseModel):
+ """One active or queued scenario run in scheduler order."""
+
+ scenario_result_id: str
+ scenario_name: str
+ scenario_registry_name: str
+ created_at: AwareDatetime
+ enqueued_at: AwareDatetime
+ started_at: AwareDatetime | None = None
+ state: ScenarioRunState
+ position: int | None = Field(None, ge=1)
+
+
+class ScenarioQueueSnapshot(BaseModel):
+ """Point-in-time FIFO scheduler state."""
+
+ revision: int = Field(ge=0)
+ snapshot_at: AwareDatetime
+ active: ScenarioQueueEntry | None = None
+ queued: list[ScenarioQueueEntry] = Field(default_factory=list)
+
+
class ScenarioAttackResultDelta(BaseModel):
"""Lightweight memory projection used to map one scenario progress delta."""
diff --git a/pyrit/registry/components/scenario_registry.py b/pyrit/registry/components/scenario_registry.py
index 7e8af09f46..5be7fea7dd 100644
--- a/pyrit/registry/components/scenario_registry.py
+++ b/pyrit/registry/components/scenario_registry.py
@@ -24,6 +24,7 @@
from pyrit.registry.registry_metadata import RegistryMetadata
if TYPE_CHECKING:
+ from collections.abc import Mapping
from types import ModuleType
from pyrit.models import Parameter
@@ -226,6 +227,7 @@ async def create_and_initialize_async(
*,
scenario_params: dict[str, Any] | None = None,
scenario_result_id: str | None = None,
+ initial_metadata: Mapping[str, Any] | None = None,
**initialize_kwargs: Any,
) -> Scenario:
"""
@@ -255,6 +257,8 @@ async def create_and_initialize_async(
parameters to set before initialization. Defaults to an empty mapping.
scenario_result_id (str | None): Existing scenario-result id to resume,
or ``None`` to start a fresh run.
+ initial_metadata (Mapping[str, Any] | None): Caller-owned metadata to
+ persist atomically when a fresh scenario result is created.
**initialize_kwargs (Any): Common run-resolved parameters merged into the
param bag (notably ``objective_target``).
@@ -268,5 +272,7 @@ async def create_and_initialize_async(
merged_args = {**(scenario_params or {}), **initialize_kwargs}
scenario = self._create_and_configure(name, params=merged_args, constructor_kwargs=constructor_kwargs)
scenario.set_scenario_registry_name(scenario_registry_name=name)
+ if initial_metadata:
+ scenario.set_initial_metadata(metadata=initial_metadata)
await scenario.initialize_async()
return scenario
diff --git a/pyrit/scenario/core/scenario.py b/pyrit/scenario/core/scenario.py
index 75f236e627..0f94633de7 100644
--- a/pyrit/scenario/core/scenario.py
+++ b/pyrit/scenario/core/scenario.py
@@ -12,7 +12,7 @@
import logging
import uuid
from abc import ABC, abstractmethod
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
from enum import Enum
from pathlib import Path
from typing import TYPE_CHECKING, Any, ClassVar, Literal, final
@@ -245,6 +245,7 @@ def __init__(
self._atomic_attacks: list[AtomicAttack] = []
self._scenario_result_id: str | None = str(scenario_result_id) if scenario_result_id else None
self._scenario_registry_name: str | None = None
+ self._initial_metadata: dict[str, Any] = {}
self._active_atomic_groups: dict[str, str] = {}
# Store prepared techniques for use in _build_atomic_attacks_async
@@ -305,6 +306,10 @@ def set_scenario_registry_name(self, *, scenario_registry_name: str) -> None:
"""Record the requested registry name for durable run-plan attribution."""
self._scenario_registry_name = scenario_registry_name
+ def set_initial_metadata(self, *, metadata: Mapping[str, Any]) -> None:
+ """Set caller-owned metadata to persist when a new scenario result is created."""
+ self._initial_metadata = dict(metadata)
+
@classmethod
def _common_scenario_parameters(cls) -> list[Parameter]:
"""
@@ -923,7 +928,10 @@ async def initialize_async(self) -> None:
attack_results=attack_results,
scenario_run_state=ScenarioRunState.CREATED,
display_group_map=self._display_group_map,
- metadata=self._build_initial_scenario_metadata(),
+ metadata={
+ **self._build_initial_scenario_metadata(),
+ **self._initial_metadata,
+ },
)
self._memory.add_scenario_results_to_memory(scenario_results=[result])
diff --git a/tests/unit/backend/test_main.py b/tests/unit/backend/test_main.py
index 19e1471e2a..dc70945243 100644
--- a/tests/unit/backend/test_main.py
+++ b/tests/unit/backend/test_main.py
@@ -18,13 +18,26 @@
from starlette.exceptions import HTTPException as StarletteHTTPException
from pyrit.backend.main import SPAStaticFiles, app, lifespan, setup_frontend
+from pyrit.backend.services.scenario_run_service import ScenarioRunService
+from pyrit.memory import AzureSQLMemory
from pyrit.setup.configuration_loader import ConfigurationLoader
+@pytest.fixture
+def mock_scenario_run_lifecycle():
+ """Mock scenario scheduling lifecycle hooks."""
+ service = MagicMock(
+ reconcile_interrupted_runs_async=AsyncMock(return_value=0),
+ shutdown_async=AsyncMock(),
+ )
+ with patch("pyrit.backend.main.get_scenario_run_service", return_value=service):
+ yield service
+
+
class TestLifespan:
"""Tests for the application lifespan context manager."""
- async def test_lifespan_yields(self) -> None:
+ async def test_lifespan_yields(self, mock_scenario_run_lifecycle) -> None:
"""Test that lifespan delegates to ConfigurationLoader and yields."""
fake_config = ConfigurationLoader()
with (
@@ -43,8 +56,10 @@ async def test_lifespan_yields(self) -> None:
assert app.state.default_labels == {}
assert app.state.max_concurrent_scenario_runs == fake_config.max_concurrent_scenario_runs
assert app.state.allow_custom_initializers is False
+ mock_scenario_run_lifecycle.reconcile_interrupted_runs_async.assert_awaited_once()
+ mock_scenario_run_lifecycle.shutdown_async.assert_awaited_once()
- async def test_lifespan_warns_when_custom_initializers_allowed(self) -> None:
+ async def test_lifespan_warns_when_custom_initializers_allowed(self, mock_scenario_run_lifecycle) -> None:
"""Test that lifespan logs a warning when allow_custom_initializers is enabled."""
fake_config = ConfigurationLoader(allow_custom_initializers=True)
with (
@@ -62,7 +77,34 @@ async def test_lifespan_warns_when_custom_initializers_allowed(self) -> None:
mock_warning.assert_called_once()
- async def test_lifespan_populates_default_labels_from_operator_and_operation(self) -> None:
+ async def test_lifespan_shared_memory_reconciliation_is_non_destructive(self) -> None:
+ shared_memory = MagicMock(spec=AzureSQLMemory)
+ fake_config = ConfigurationLoader()
+ with patch(
+ "pyrit.backend.services.scenario_run_service.CentralMemory.get_memory_instance",
+ return_value=shared_memory,
+ ):
+ service = ScenarioRunService()
+
+ with (
+ patch.object(ConfigurationLoader, "load_with_overrides", return_value=fake_config),
+ patch.object(ConfigurationLoader, "initialize_pyrit_async", new=AsyncMock()),
+ patch(
+ "pyrit.backend.main.get_initializer_service",
+ return_value=MagicMock(run_additional_initializers_async=AsyncMock()),
+ ),
+ patch("pyrit.backend.main.get_scenario_run_service", return_value=service),
+ patch("pyrit.backend.main.setup_frontend"),
+ ):
+ async with lifespan(app):
+ pass
+
+ shared_memory.get_scenario_run_state_page.assert_not_called()
+ shared_memory.update_scenario_run_state.assert_not_called()
+
+ async def test_lifespan_populates_default_labels_from_operator_and_operation(
+ self, mock_scenario_run_lifecycle
+ ) -> None:
"""Test that operator and operation are exposed as default_labels."""
fake_config = ConfigurationLoader(operator="alice", operation="op-42")
with (
@@ -79,7 +121,7 @@ async def test_lifespan_populates_default_labels_from_operator_and_operation(sel
assert app.state.default_labels == {"operator": "alice", "operation": "op-42"}
- async def test_lifespan_reads_config_file_env_var(self) -> None:
+ async def test_lifespan_reads_config_file_env_var(self, mock_scenario_run_lifecycle) -> None:
"""Test that PYRIT_CONFIG_FILE is forwarded to ConfigurationLoader.load_with_overrides."""
fake_config = ConfigurationLoader()
with (
diff --git a/tests/unit/backend/test_scenario_run_routes.py b/tests/unit/backend/test_scenario_run_routes.py
index 1a6b5b4f00..999af585e0 100644
--- a/tests/unit/backend/test_scenario_run_routes.py
+++ b/tests/unit/backend/test_scenario_run_routes.py
@@ -22,11 +22,13 @@
AttackOutcome,
AttackResult,
ScenarioProgressHeader,
+ ScenarioQueueEntry,
+ ScenarioQueueSnapshot,
ScenarioRunPlan,
ScenarioRunProgress,
ScenarioRunState,
)
-from pyrit.models.catalog.scenario import ScenarioRunSummary
+from pyrit.models.catalog import ScenarioRunSummary
from unit.mocks import make_scenario_result
@@ -238,6 +240,45 @@ def test_list_runs_returns_400_for_invalid_cursor(self, client: TestClient) -> N
assert response.json()["detail"] == "Malformed scenario history cursor."
+class TestScenarioRunQueueRoute:
+ """Tests for GET /api/scenarios/runs/queue."""
+
+ def test_queue_returns_active_and_ordered_entries(self, client: TestClient) -> None:
+ now = datetime(2025, 1, 1, tzinfo=timezone.utc)
+ snapshot = ScenarioQueueSnapshot(
+ revision=4,
+ snapshot_at=now,
+ active=ScenarioQueueEntry(
+ scenario_result_id="active",
+ scenario_name="ActiveScenario",
+ scenario_registry_name="active.scenario",
+ state=ScenarioRunState.IN_PROGRESS,
+ created_at=now,
+ enqueued_at=now,
+ started_at=now,
+ ),
+ queued=[
+ ScenarioQueueEntry(
+ scenario_result_id="queued",
+ scenario_name="QueuedScenario",
+ scenario_registry_name="queued.scenario",
+ state=ScenarioRunState.QUEUED,
+ position=1,
+ created_at=now,
+ enqueued_at=now,
+ )
+ ],
+ )
+ with patch("pyrit.backend.routes.scenarios.get_scenario_run_service") as mock_get:
+ mock_get.return_value.get_queue_snapshot.return_value = snapshot
+
+ response = client.get("/api/scenarios/runs/queue")
+
+ assert response.status_code == status.HTTP_200_OK
+ assert response.json()["active"]["scenario_result_id"] == "active"
+ assert response.json()["queued"][0]["position"] == 1
+
+
class TestGetScenarioRunRoute:
"""Tests for GET /api/scenarios/runs/{id}."""
@@ -337,7 +378,12 @@ def test_progress_returns_compact_plan_response(self, client: TestClient) -> Non
with patch("pyrit.backend.routes.scenarios.get_scenario_run_service") as mock_get:
mock_service = MagicMock()
mock_service.snapshot_active_run.side_effect = lambda **_: (
- snapshot_thread.append(get_ident()) or MagicMock(active_group_ids=("active-group",))
+ snapshot_thread.append(get_ident())
+ or MagicMock(
+ active_group_ids=("active-group",),
+ queue_position=None,
+ active_scenario_result_id="test-run-id",
+ )
)
mock_service.get_run_progress_from_storage.side_effect = lambda **_: (
storage_thread.append(get_ident()) or progress
@@ -354,6 +400,8 @@ def test_progress_returns_compact_plan_response(self, client: TestClient) -> Non
since=None,
limit=25,
active_group_ids=("active-group",),
+ queue_position=None,
+ active_scenario_result_id="test-run-id",
)
assert snapshot_thread[0] != storage_thread[0]
@@ -376,7 +424,11 @@ async def test_progress_supports_direct_keyword_call(self) -> None:
)
with patch("pyrit.backend.routes.scenarios.get_scenario_run_service") as mock_get:
mock_service = MagicMock()
- mock_service.snapshot_active_run.return_value = MagicMock(active_group_ids=())
+ mock_service.snapshot_active_run.return_value = MagicMock(
+ active_group_ids=(),
+ queue_position=None,
+ active_scenario_result_id="test-run-id",
+ )
mock_service.get_run_progress_from_storage.return_value = progress
mock_get.return_value = mock_service
@@ -392,6 +444,8 @@ async def test_progress_supports_direct_keyword_call(self) -> None:
since=None,
limit=25,
active_group_ids=(),
+ queue_position=None,
+ active_scenario_result_id="test-run-id",
)
diff --git a/tests/unit/backend/test_scenario_run_service.py b/tests/unit/backend/test_scenario_run_service.py
index 19022c6358..47569bccbd 100644
--- a/tests/unit/backend/test_scenario_run_service.py
+++ b/tests/unit/backend/test_scenario_run_service.py
@@ -16,11 +16,16 @@
import pyrit.backend.services.scenario_run_service as _svc_mod
from pyrit.backend.services.scenario_run_service import (
- _DEFAULT_MAX_CONCURRENT_RUNS,
ScenarioRunService,
)
from pyrit.converter import Converter
-from pyrit.memory import ScenarioHistoryRunRecord, ScenarioHistoryUnitRecord
+from pyrit.memory import (
+ AzureSQLMemory,
+ ScenarioHistoryRunRecord,
+ ScenarioHistoryUnitRecord,
+ ScenarioRunStateRecord,
+ SQLiteMemory,
+)
from pyrit.models import (
SCENARIO_RUN_PLAN_METADATA_KEY,
AtomicAttackIdentifier,
@@ -28,6 +33,7 @@
AttackResult,
AttackSeedGroup,
ComponentIdentifier,
+ RetryEvent,
ScenarioAttackResultDelta,
ScenarioResult,
ScenarioRunPlan,
@@ -149,6 +155,7 @@ def _make_history_record(
status=run_state.value,
labels={},
created_at=scenario_result.creation_time,
+ started_at=None,
completed_at=scenario_result.completion_time,
error_message=None,
error_type=None,
@@ -161,7 +168,7 @@ def _make_history_record(
@pytest.fixture
def mock_memory():
"""Patch CentralMemory.get_memory_instance to return a mock."""
- mock = MagicMock()
+ mock = MagicMock(spec=SQLiteMemory)
mock.get_scenario_results.return_value = []
# Default: no error AttackResults linked to any scenario. Tests that exercise
# the error fallback path explicitly set get_attack_results.return_value.
@@ -220,12 +227,28 @@ class TestScenarioRunServiceStartRun:
async def test_start_run_returns_running_status(self, mock_all_registries) -> None:
"""Test that starting a run returns RUNNING status with run_id = scenario_result_id."""
service = ScenarioRunService()
+ mock_memory = mock_all_registries["memory"]
+ service._terminal_errors["sr-uuid-1"] = "prior failed attempt"
response = await service.start_run_async(request=_make_request())
assert response.scenario_result_id == "sr-uuid-1"
assert response.status == ScenarioRunState.IN_PROGRESS
assert response.scenario_name == "foundry.red_team_agent"
+ assert "sr-uuid-1" not in service._terminal_errors
assert response.error is None
+ metadata_call = mock_memory.update_scenario_run_state_and_metadata_fields.call_args
+ assert metadata_call.kwargs["scenario_result_id"] == "sr-uuid-1"
+ assert metadata_call.kwargs["scenario_run_state"] == ScenarioRunState.IN_PROGRESS
+ persisted_start = datetime.fromisoformat(
+ metadata_call.kwargs["metadata_fields"][_svc_mod.SCENARIO_RUN_STARTED_AT_METADATA_KEY]
+ )
+ assert persisted_start.tzinfo is not None
+ assert (
+ metadata_call.kwargs["metadata_fields"][_svc_mod._SCHEDULER_METADATA_KEY]
+ == _svc_mod._SCHEDULER_METADATA_VALUE
+ )
+ mock_memory.update_scenario_metadata_fields.assert_not_called()
+ mock_memory.update_scenario_run_state.assert_not_called()
async def test_start_run_invalid_scenario_raises_value_error(self, mock_memory) -> None:
"""Test that an invalid scenario name raises ValueError immediately."""
@@ -374,6 +397,7 @@ def get_aggregate_tags(cls) -> set[str]:
"airt.jailbreak",
scenario_params=scenario_params,
scenario_result_id=None,
+ initial_metadata={_svc_mod._SCHEDULER_METADATA_KEY: _svc_mod._SCHEDULER_METADATA_VALUE},
objective_target=objective_target,
max_concurrency=10,
max_retries=0,
@@ -381,6 +405,107 @@ def get_aggregate_tags(cls) -> set[str]:
scenario_techniques=[_JailbreakTechnique.PROMPT_SENDING],
)
+ async def test_exact_eight_unit_jailbreak_request_queues_behind_active_run(self, mock_all_registries) -> None:
+ """The configured eight-unit request keeps a stable ID and FIFO position while another run executes."""
+
+ class _JailbreakTechnique(ScenarioTechnique):
+ ALL = ("all", {"all"})
+ DEFAULT = ("default", {"default"})
+ PROMPT_SENDING = ("prompt_sending", {"default"})
+
+ @classmethod
+ def get_aggregate_tags(cls) -> set[str]:
+ return {"all", "default"}
+
+ service = ScenarioRunService()
+ mock_sr = mock_all_registries["scenario_registry"]
+ mock_memory = mock_all_registries["memory"]
+ mock_all_registries["scenario_instance"]._technique_class = _JailbreakTechnique
+ records: dict[str, MagicMock] = {}
+ active_started = asyncio.Event()
+ queued_started = asyncio.Event()
+ release_active = asyncio.Event()
+ started: list[str] = []
+
+ async def _create_scenario(*args: object, **kwargs: object) -> MagicMock:
+ run_id = f"run-{len(records) + 1}"
+ record = _make_db_scenario_result(
+ result_id=run_id,
+ scenario_name=str(args[0]),
+ run_state=ScenarioRunState.CREATED,
+ )
+ records[run_id] = record
+ scenario = MagicMock()
+ scenario._scenario_result_id = run_id
+ scenario.active_atomic_group_ids = set()
+
+ async def _run() -> None:
+ started.append(run_id)
+ if run_id == "run-1":
+ active_started.set()
+ await release_active.wait()
+ else:
+ queued_started.set()
+ record.scenario_run_state = ScenarioRunState.COMPLETED
+
+ scenario.run_async = AsyncMock(side_effect=_run)
+ return scenario
+
+ def _get_results(*, scenario_result_ids: list[str] | None = None) -> list[MagicMock]:
+ if scenario_result_ids is None:
+ return list(records.values())
+ return [records[run_id] for run_id in scenario_result_ids if run_id in records]
+
+ def _update_state(*, scenario_result_id: str, scenario_run_state: ScenarioRunState, **_: object) -> None:
+ records[scenario_result_id].scenario_run_state = scenario_run_state
+
+ mock_sr.create_and_initialize_async = AsyncMock(side_effect=_create_scenario)
+ mock_memory.get_scenario_results.side_effect = _get_results
+ mock_memory.update_scenario_run_state.side_effect = _update_state
+ mock_memory.update_scenario_run_state_and_metadata_fields.side_effect = _update_state
+
+ active_response = await service.start_run_async(request=_make_request())
+ await asyncio.wait_for(active_started.wait(), timeout=1)
+ configured_request = _make_request(
+ scenario_name="airt.jailbreak",
+ techniques=["prompt_sending"],
+ include_baseline=False,
+ scenario_params={"num_jailbreaks": 2, "num_jailbreak_attempts": 1},
+ )
+ queued_response = await service.start_run_async(request=configured_request)
+
+ assert active_response.scenario_result_id == "run-1"
+ assert queued_response.scenario_result_id == "run-2"
+ assert queued_response.status == ScenarioRunState.QUEUED
+ assert queued_response.queue_position == 1
+ assert queued_response.active_scenario_result_id == "run-1"
+ queued_transition = next(
+ call
+ for call in mock_memory.update_scenario_run_state_and_metadata_fields.call_args_list
+ if call.kwargs["scenario_result_id"] == "run-2"
+ and call.kwargs["scenario_run_state"] == ScenarioRunState.QUEUED
+ )
+ assert (
+ queued_transition.kwargs["metadata_fields"][_svc_mod._SCHEDULER_METADATA_KEY]
+ == _svc_mod._SCHEDULER_METADATA_VALUE
+ )
+ assert [(entry.scenario_result_id, entry.position) for entry in service.get_queue_snapshot().queued] == [
+ ("run-2", 1)
+ ]
+ second_init = mock_sr.create_and_initialize_async.await_args_list[1]
+ assert second_init.args == ("airt.jailbreak",)
+ assert second_init.kwargs["scenario_params"] == {
+ "num_jailbreaks": 2,
+ "num_jailbreak_attempts": 1,
+ }
+ assert second_init.kwargs["scenario_techniques"] == [_JailbreakTechnique.PROMPT_SENDING]
+ assert second_init.kwargs["include_baseline"] is False
+
+ release_active.set()
+ await asyncio.wait_for(queued_started.wait(), timeout=1)
+ await asyncio.wait_for(service._active_tasks["run-2"].task, timeout=1)
+ assert started == ["run-1", "run-2"]
+
async def test_start_run_forwards_include_baseline(self, mock_all_registries) -> None:
service = ScenarioRunService()
request = _make_request()
@@ -569,30 +694,112 @@ class _MarkerDatasetConfiguration(DatasetConfiguration):
assert built_config.dataset_names == ["a", "b"]
assert built_config.max_dataset_size == 7
- async def test_start_run_exceeds_concurrent_limit(self, mock_all_registries) -> None:
- """Test that exceeding concurrent run limit raises ValueError."""
+ async def test_concurrent_launches_run_one_at_a_time_in_fifo_order(self, mock_all_registries) -> None:
+ """Concurrent launches queue durably and hand off exactly once in FIFO order."""
service = ScenarioRunService()
- scenario_instance = mock_all_registries["scenario_instance"]
mock_sr = mock_all_registries["scenario_registry"]
+ mock_memory = mock_all_registries["memory"]
+ records: dict[str, MagicMock] = {}
+ release_events: dict[str, asyncio.Event] = {}
+ started_events: dict[str, asyncio.Event] = {}
+ started: list[str] = []
+ active_count = 0
+ max_active_count = 0
+ fail_once = {"handoff_read": False, "queued_cancel": False, "active_cancel": False}
+
+ async def _create_scenario(*args: object, **kwargs: object) -> MagicMock:
+ run_id = f"run-{len(records) + 1}"
+ record = _make_db_scenario_result(result_id=run_id, run_state=ScenarioRunState.CREATED)
+ records[run_id] = record
+ release_events[run_id] = asyncio.Event()
+ started_events[run_id] = asyncio.Event()
+ scenario = MagicMock()
+ scenario._scenario_result_id = run_id
+
+ async def _run() -> None:
+ nonlocal active_count, max_active_count
+ active_count += 1
+ max_active_count = max(max_active_count, active_count)
+ started.append(run_id)
+ started_events[run_id].set()
+ try:
+ await release_events[run_id].wait()
+ except asyncio.CancelledError:
+ raise
+ else:
+ record.scenario_run_state = ScenarioRunState.COMPLETED
+ finally:
+ active_count -= 1
+
+ scenario.run_async = AsyncMock(side_effect=_run)
+ return scenario
+
+ def _get_results(*, scenario_result_ids: list[str] | None = None) -> list[MagicMock]:
+ if scenario_result_ids is None:
+ return list(records.values())
+ if fail_once["handoff_read"] and scenario_result_ids == ["run-2"]:
+ fail_once["handoff_read"] = False
+ raise RuntimeError("temporary storage failure")
+ return [records[run_id] for run_id in scenario_result_ids if run_id in records]
+
+ def _update_state(*, scenario_result_id: str, scenario_run_state: ScenarioRunState, **_: object) -> None:
+ if (
+ fail_once["queued_cancel"]
+ and scenario_result_id == "run-3"
+ and scenario_run_state == ScenarioRunState.CANCELLED
+ ):
+ fail_once["queued_cancel"] = False
+ raise RuntimeError("temporary cancellation persistence failure")
+ if (
+ fail_once["active_cancel"]
+ and scenario_result_id == "run-2"
+ and scenario_run_state == ScenarioRunState.CANCELLED
+ ):
+ fail_once["active_cancel"] = False
+ raise RuntimeError("temporary active cancellation persistence failure")
+ records[scenario_result_id].scenario_run_state = scenario_run_state
+
+ mock_sr.create_and_initialize_async = AsyncMock(side_effect=_create_scenario)
+ mock_memory.get_scenario_results.side_effect = _get_results
+ mock_memory.update_scenario_run_state.side_effect = _update_state
+ mock_memory.update_scenario_run_state_and_metadata_fields.side_effect = _update_state
+
+ responses = await asyncio.gather(*(service.start_run_async(request=_make_request()) for _ in range(4)))
+
+ assert [response.scenario_result_id for response in responses] == ["run-1", "run-2", "run-3", "run-4"]
+ snapshot = service.get_queue_snapshot()
+ assert snapshot.active and snapshot.active.scenario_result_id == "run-1"
+ assert [(entry.scenario_result_id, entry.position) for entry in snapshot.queued] == [
+ ("run-2", 1),
+ ("run-3", 2),
+ ("run-4", 3),
+ ]
- # Each call needs a unique scenario_result_id
- call_count = 0
-
- async def _set_unique_id(*args: object, **kwargs: object) -> object:
- nonlocal call_count
- call_count += 1
- scenario_instance._scenario_result_id = f"sr-uuid-{call_count}"
- return scenario_instance
-
- mock_sr.create_and_initialize_async = AsyncMock(side_effect=_set_unique_id)
+ fail_once["handoff_read"] = True
+ release_events["run-1"].set()
+ await asyncio.wait_for(started_events["run-2"].wait(), timeout=1)
+ fail_once["queued_cancel"] = True
+ with pytest.raises(RuntimeError, match="temporary cancellation persistence failure"):
+ await service.cancel_run_async(scenario_result_id="run-3")
+ assert [entry.scenario_result_id for entry in service.get_queue_snapshot().queued] == ["run-3", "run-4"]
+ cancelled = await service.cancel_run_async(scenario_result_id="run-3")
+ assert cancelled and cancelled.status == ScenarioRunState.CANCELLED
+ assert [(entry.scenario_result_id, entry.position) for entry in service.get_queue_snapshot().queued] == [
+ ("run-4", 1)
+ ]
- # Fill up to the limit
- for _ in range(_DEFAULT_MAX_CONCURRENT_RUNS):
- await service.start_run_async(request=_make_request())
+ fail_once["active_cancel"] = True
+ with pytest.raises(RuntimeError, match="temporary active cancellation persistence failure"):
+ await service.cancel_run_async(scenario_result_id="run-2")
+ await asyncio.wait_for(started_events["run-4"].wait(), timeout=1)
+ assert records["run-2"].scenario_run_state == ScenarioRunState.CANCELLED
+ release_events["run-4"].set()
+ await asyncio.wait_for(service._active_tasks["run-4"].task, timeout=1)
- # Next one should fail
- with pytest.raises(ValueError, match="Maximum concurrent runs"):
- await service.start_run_async(request=_make_request())
+ assert started == ["run-1", "run-2", "run-4"]
+ assert max_active_count == 1
+ assert service.get_queue_snapshot().active is None
+ assert service.get_queue_snapshot().queued == []
async def test_start_run_runs_initializers(self, mock_all_registries) -> None:
"""Test that initializers are run during start_run_async."""
@@ -986,6 +1193,7 @@ async def test_cancel_run_sets_cancelled_status(self, mock_all_registries) -> No
"""Test that cancelling a running scenario persists CANCELLED to DB."""
service = ScenarioRunService()
mock_memory = mock_all_registries["memory"]
+ mock_all_registries["scenario_instance"].run_async.side_effect = asyncio.Event().wait
response = await service.start_run_async(request=_make_request())
# After update_scenario_run_state, the next DB query should return CANCELLED
@@ -998,7 +1206,7 @@ async def test_cancel_run_sets_cancelled_status(self, mock_all_registries) -> No
result = await service.cancel_run_async(scenario_result_id=response.scenario_result_id)
- mock_memory.update_scenario_run_state.assert_called_once_with(
+ mock_memory.update_scenario_run_state.assert_any_call(
scenario_result_id=response.scenario_result_id,
scenario_run_state=ScenarioRunState.CANCELLED,
error_message="Run was cancelled by user",
@@ -1007,6 +1215,157 @@ async def test_cancel_run_sets_cancelled_status(self, mock_all_registries) -> No
assert result is not None
assert result.status == ScenarioRunState.CANCELLED
+
+class TestScenarioRunServiceRecovery:
+ """Tests for restart reconciliation and overload evidence."""
+
+ async def test_reconcile_marks_only_scheduler_managed_local_rows_failed(self, mock_memory) -> None:
+ scheduler_metadata = {_svc_mod._SCHEDULER_METADATA_KEY: _svc_mod._SCHEDULER_METADATA_VALUE}
+ interrupted = [
+ ScenarioRunStateRecord(
+ scenario_result_id="created",
+ state=ScenarioRunState.CREATED,
+ ),
+ ScenarioRunStateRecord(
+ scenario_result_id="queued",
+ state=ScenarioRunState.QUEUED,
+ ),
+ ScenarioRunStateRecord(
+ scenario_result_id="running",
+ state=ScenarioRunState.IN_PROGRESS,
+ ),
+ ScenarioRunStateRecord(
+ scenario_result_id="framework-run",
+ state=ScenarioRunState.IN_PROGRESS,
+ ),
+ ]
+ mock_memory.get_scenario_run_state_page.return_value = (interrupted, False)
+ headers = {
+ "created": MagicMock(metadata=scheduler_metadata),
+ "queued": MagicMock(metadata=scheduler_metadata),
+ "running": MagicMock(metadata=scheduler_metadata),
+ "framework-run": MagicMock(metadata={}),
+ }
+ mock_memory.get_scenario_result_header.side_effect = lambda *, scenario_result_id: headers[scenario_result_id]
+
+ reconciled = await ScenarioRunService().reconcile_interrupted_runs_async()
+
+ assert reconciled == 3
+ assert {call.kwargs["scenario_result_id"] for call in mock_memory.update_scenario_run_state.call_args_list} == {
+ "created",
+ "queued",
+ "running",
+ }
+ assert all(
+ call.kwargs["scenario_run_state"] == ScenarioRunState.FAILED
+ and call.kwargs["error_type"] == "ScenarioInterruptedError"
+ for call in mock_memory.update_scenario_run_state.call_args_list
+ )
+ mock_memory.get_scenario_results.assert_not_called()
+ mock_memory.get_scenario_run_state_page.assert_called_once_with(
+ states=(ScenarioRunState.CREATED, ScenarioRunState.QUEUED, ScenarioRunState.IN_PROGRESS),
+ after_id=None,
+ limit=500,
+ )
+
+ async def test_reconcile_pages_nonterminal_state_projection(self, mock_memory) -> None:
+ first = ScenarioRunStateRecord(
+ scenario_result_id="00000000-0000-0000-0000-000000000001",
+ state=ScenarioRunState.QUEUED,
+ )
+ second = ScenarioRunStateRecord(
+ scenario_result_id="00000000-0000-0000-0000-000000000002",
+ state=ScenarioRunState.IN_PROGRESS,
+ )
+ mock_memory.get_scenario_run_state_page.side_effect = [([first], True), ([second], False)]
+ mock_memory.get_scenario_result_header.return_value = MagicMock(
+ metadata={_svc_mod._SCHEDULER_METADATA_KEY: _svc_mod._SCHEDULER_METADATA_VALUE}
+ )
+
+ reconciled = await ScenarioRunService().reconcile_interrupted_runs_async()
+
+ assert reconciled == 2
+ assert mock_memory.get_scenario_run_state_page.call_args_list[1].kwargs["after_id"] == first.scenario_result_id
+
+ async def test_reconcile_shared_backend_is_non_destructive(self) -> None:
+ shared_memory = MagicMock(spec=AzureSQLMemory)
+ with patch(_MEMORY_PATCH, return_value=shared_memory):
+ reconciled = await ScenarioRunService().reconcile_interrupted_runs_async()
+
+ assert reconciled == 0
+ shared_memory.get_scenario_run_state_page.assert_not_called()
+ shared_memory.update_scenario_run_state.assert_not_called()
+
+ async def test_shutdown_fails_active_and_queued_runs_without_starting_next(self, mock_all_registries) -> None:
+ mock_scenario_registry = mock_all_registries["scenario_registry"]
+ mock_memory = mock_all_registries["memory"]
+ records: dict[str, MagicMock] = {}
+ scenarios: dict[str, MagicMock] = {}
+
+ async def _create_scenario(*args: object, **kwargs: object) -> MagicMock:
+ run_id = f"shutdown-{len(records) + 1}"
+ records[run_id] = _make_db_scenario_result(
+ result_id=run_id,
+ run_state=ScenarioRunState.CREATED,
+ )
+ scenario = MagicMock()
+ scenario._scenario_result_id = run_id
+ scenario.run_async = AsyncMock(side_effect=asyncio.Event().wait)
+ scenarios[run_id] = scenario
+ return scenario
+
+ def _get_results(*, scenario_result_ids: list[str] | None = None) -> list[MagicMock]:
+ if scenario_result_ids is None:
+ return list(records.values())
+ return [records[run_id] for run_id in scenario_result_ids if run_id in records]
+
+ def _update_state(*, scenario_result_id: str, scenario_run_state: ScenarioRunState, **_: object) -> None:
+ records[scenario_result_id].scenario_run_state = scenario_run_state
+
+ mock_scenario_registry.create_and_initialize_async = AsyncMock(side_effect=_create_scenario)
+ mock_memory.get_scenario_results.side_effect = _get_results
+ mock_memory.update_scenario_run_state.side_effect = _update_state
+ mock_memory.update_scenario_run_state_and_metadata_fields.side_effect = _update_state
+ service = ScenarioRunService()
+ await service.start_run_async(request=_make_request())
+ await service.start_run_async(request=_make_request())
+ await asyncio.sleep(0)
+
+ await service.shutdown_async()
+
+ assert records["shutdown-1"].scenario_run_state == ScenarioRunState.FAILED
+ assert records["shutdown-2"].scenario_run_state == ScenarioRunState.FAILED
+ scenarios["shutdown-1"].run_async.assert_awaited_once()
+ scenarios["shutdown-2"].run_async.assert_not_awaited()
+ failure_calls = [
+ call
+ for call in mock_memory.update_scenario_run_state.call_args_list
+ if call.kwargs.get("scenario_run_state") == ScenarioRunState.FAILED
+ ]
+ assert len(failure_calls) == 2
+ assert all(call.kwargs["error_type"] == "ScenarioInterruptedError" for call in failure_calls)
+ assert all("shut down" in call.kwargs["error_message"] for call in failure_calls)
+
+ def test_overload_summaries_group_429_and_5xx_by_role_without_false_positives(self, mock_memory) -> None:
+ now = datetime(2025, 1, 1, tzinfo=timezone.utc)
+ events = [
+ RetryEvent(component_role="adversarial_chat", status_code=429, timestamp=now),
+ RetryEvent(component_role="adversarial_chat", status_code=503, timestamp=now + timedelta(seconds=2)),
+ RetryEvent(component_role="objective_target", status_code=500, timestamp=now + timedelta(seconds=1)),
+ RetryEvent(component_role="objective_target", status_code=408, timestamp=now + timedelta(seconds=3)),
+ RetryEvent(component_role="objective_target", exception_message="HTTP 429", timestamp=now),
+ ]
+
+ summaries = ScenarioRunService._build_overload_summaries(retry_events=events)
+
+ assert [summary.component_role for summary in summaries] == ["adversarial_chat", "objective_target"]
+ assert summaries[0].count == 2
+ assert summaries[0].rate_limit_count == 1
+ assert summaries[0].server_error_count == 1
+ assert summaries[0].status_codes == [429, 503]
+ assert summaries[1].count == 1
+ assert summaries[1].status_codes == [500]
+
async def test_cancel_waits_for_final_persisted_progress_delta(self, mock_all_registries) -> None:
"""Cancellation completes task cleanup before callers can fetch terminal progress."""
mock_memory = mock_all_registries["memory"]
@@ -1090,46 +1449,62 @@ async def test_execute_run_completes_successfully(self, mock_all_registries) ->
mock_scenario_result.creation_time = datetime(2025, 1, 1, tzinfo=timezone.utc)
mock_scenario_result.completion_time = datetime(2025, 1, 1, 0, 5, tzinfo=timezone.utc)
- mock_instance.run_async = AsyncMock(return_value=mock_scenario_result)
+ execution_started = asyncio.Event()
+ release_execution = asyncio.Event()
+
+ async def _run() -> MagicMock:
+ execution_started.set()
+ await release_execution.wait()
+ return mock_scenario_result
+
+ mock_instance.run_async = AsyncMock(side_effect=_run)
response = await service.start_run_async(request=_make_request())
+ await execution_started.wait()
# Wait for the background task to complete
active = service._active_tasks.get(response.scenario_result_id)
assert active is not None
assert active.task is not None
+ release_execution.set()
await active.task
- # Active task is cleaned up on next get_run (deferred cleanup)
- assert response.scenario_result_id in service._active_tasks
+ # Executable task state is released during terminal handoff.
+ assert response.scenario_result_id not in service._active_tasks
fetched = service.get_run(scenario_result_id=response.scenario_result_id)
assert fetched is not None
- assert response.scenario_result_id not in service._active_tasks
async def test_execute_run_fails_with_error(self, mock_all_registries) -> None:
"""Test that a run_async failure stores error and surfaces it via get_run."""
service = ScenarioRunService()
mock_instance = mock_all_registries["scenario_instance"]
+ execution_started = asyncio.Event()
+ release_execution = asyncio.Event()
- mock_instance.run_async = AsyncMock(side_effect=RuntimeError("scenario exploded"))
+ async def _run() -> None:
+ execution_started.set()
+ await release_execution.wait()
+ raise RuntimeError("scenario exploded")
+ mock_instance.run_async = AsyncMock(side_effect=_run)
response = await service.start_run_async(request=_make_request())
+ await execution_started.wait()
# Wait for the background task
active = service._active_tasks.get(response.scenario_result_id)
assert active is not None
assert active.task is not None
+ release_execution.set()
await active.task
- # Error is stored on the active task until get_run reads it
+ # Error evidence remains available after executable task state is released.
assert active.error == "scenario exploded"
- assert response.scenario_result_id in service._active_tasks
+ assert response.scenario_result_id not in service._active_tasks
- # get_run should surface the error and clean up
+ # get_run surfaces the bounded terminal error evidence.
fetched = service.get_run(scenario_result_id=response.scenario_result_id)
assert fetched is not None
assert fetched.error == "scenario exploded"
- assert response.scenario_result_id not in service._active_tasks
class TestScenarioRunServiceGetResults:
@@ -1294,6 +1669,29 @@ def test_error_attacks_and_retries_are_surfaced(self, mock_memory) -> None:
assert failed.error_message == "429 Too Many Requests"
assert failed.total_retries == 4
+ def test_negative_error_attack_retries_are_clamped(self, mock_memory) -> None:
+ from pyrit.models import AttackOutcome
+
+ errored = MagicMock()
+ errored.outcome = AttackOutcome.ERROR
+ errored.objective = "malformed persisted result"
+ errored.error_type = "PersistedError"
+ errored.error_message = "invalid retry count"
+ errored.total_retries = -1
+
+ db_result = _make_db_scenario_result(
+ result_id="sr-negative-retries",
+ run_state=ScenarioRunState.COMPLETED,
+ attack_results={"attack_a": [errored]},
+ )
+ mock_memory.get_scenario_results.return_value = [db_result]
+
+ fetched = ScenarioRunService().get_run(scenario_result_id="sr-negative-retries")
+
+ assert fetched is not None
+ assert fetched.total_retries == 0
+ assert fetched.failed_attacks[0].total_retries == 0
+
def test_no_failed_attacks_when_all_succeed(self, mock_memory) -> None:
from pyrit.models import AttackOutcome
@@ -1515,6 +1913,86 @@ def test_planned_progress_deduplicates_attempts_and_keeps_latest_non_error(mock_
assert summary.total_retries == 3
+def test_history_and_detail_retry_work_match_across_attempt_partitions(mock_memory) -> None:
+ objective = "partitioned objective"
+ plan = ScenarioRunPlan(
+ scenario_registry_name="test.scenario",
+ atomic_groups=[
+ ScenarioRunPlanAtomicGroup(
+ id="group-1",
+ atomic_attack_name="attack",
+ display_group="Attack",
+ technique_eval_hash="eval",
+ seed_group_ids=["seed-1"],
+ )
+ ],
+ seed_groups=[
+ ScenarioRunPlanSeedGroup(
+ id="seed-1",
+ objective_sha256=_svc_mod.to_sha256(objective),
+ objective=objective,
+ )
+ ],
+ )
+ timestamp = datetime(2026, 8, 8, tzinfo=timezone.utc)
+ attempts = [
+ AttackResult(
+ conversation_id=f"conversation-{index}",
+ objective=objective,
+ outcome=outcome,
+ total_retries=inner_retries,
+ timestamp=timestamp + timedelta(seconds=index),
+ )
+ for index, (outcome, inner_retries) in enumerate(
+ ((AttackOutcome.ERROR, 1), (AttackOutcome.ERROR, 0), (AttackOutcome.SUCCESS, 2))
+ )
+ ]
+ scenario_result = make_scenario_result(
+ attack_results={"attack": attempts},
+ scenario_run_state=ScenarioRunState.COMPLETED,
+ metadata={SCENARIO_RUN_PLAN_METADATA_KEY: plan.model_dump(mode="json")},
+ )
+ record = replace(
+ _make_history_record(result_id=str(scenario_result.id), run_state=ScenarioRunState.COMPLETED),
+ scenario_registry_name=plan.scenario_registry_name,
+ plan_atomic_groups=[group.model_dump(mode="json") for group in plan.atomic_groups],
+ plan_seed_id_map=[{"id": "seed-1", "objective_sha256": _svc_mod.to_sha256(objective)}],
+ )
+ units = [
+ ScenarioHistoryUnitRecord(
+ scenario_result_id=str(scenario_result.id),
+ atomic_attack_name="attack",
+ technique_eval_hash="eval",
+ seed_group_id=_svc_mod.to_sha256(objective),
+ objective_sha256=_svc_mod.to_sha256(objective),
+ latest_outcome=AttackOutcome.ERROR.value,
+ latest_timestamp=timestamp + timedelta(seconds=1),
+ total_retries=2,
+ error_count=2,
+ attempt_count=2,
+ ),
+ ScenarioHistoryUnitRecord(
+ scenario_result_id=str(scenario_result.id),
+ atomic_attack_name="attack",
+ technique_eval_hash="eval",
+ seed_group_id="seed-1",
+ objective_sha256=_svc_mod.to_sha256(objective),
+ latest_outcome=AttackOutcome.SUCCESS.value,
+ latest_timestamp=timestamp + timedelta(seconds=2),
+ total_retries=2,
+ error_count=0,
+ attempt_count=1,
+ ),
+ ]
+ service = ScenarioRunService()
+
+ detail = service._build_response_from_db(scenario_result=scenario_result)
+ history = service._build_history_summary(record=record, units=units)
+
+ assert detail.total_retries == 5
+ assert history.total_retries == detail.total_retries
+
+
def test_planned_progress_maps_legacy_objective_hash_to_logical_seed_id(mock_memory) -> None:
objective = "legacy resumed objective"
seed_group = AttackSeedGroup(seeds=[SeedObjective(value=objective)])
@@ -1590,6 +2068,33 @@ def test_get_progress_uses_lightweight_queries_without_full_hydration(mock_memor
assert str(header.id) not in service._active_tasks
+def test_get_progress_exposes_persisted_started_at(mock_memory) -> None:
+ started_at = datetime(2026, 8, 8, 12, 30, tzinfo=timezone.utc)
+ header = make_scenario_result(
+ attack_results={},
+ metadata={
+ SCENARIO_RUN_PLAN_METADATA_KEY: ScenarioRunPlan(
+ atomic_groups=[],
+ seed_groups=[],
+ scenario_registry_name="test.scenario",
+ ).model_dump(mode="json"),
+ _svc_mod.SCENARIO_RUN_STARTED_AT_METADATA_KEY: started_at.isoformat(),
+ },
+ )
+ mock_memory.get_scenario_result_header.return_value = header
+ mock_memory.get_scenario_attack_result_deltas.return_value = ([], False)
+
+ progress = ScenarioRunService().get_run_progress_from_storage(
+ scenario_result_id=str(header.id),
+ since=None,
+ limit=25,
+ active_group_ids=[],
+ )
+
+ assert progress is not None
+ assert progress.run.started_at == started_at
+
+
def test_get_progress_rejects_duplicate_stored_plan_groups(mock_memory) -> None:
group = ScenarioRunPlanAtomicGroup(
id="duplicate",
diff --git a/tests/unit/backend/test_scenario_service.py b/tests/unit/backend/test_scenario_service.py
index f1545a1f1e..76a03ef9c9 100644
--- a/tests/unit/backend/test_scenario_service.py
+++ b/tests/unit/backend/test_scenario_service.py
@@ -765,13 +765,16 @@ def test_get_scenario_returns_404_when_not_found(self, client: TestClient) -> No
assert response.status_code == status.HTTP_404_NOT_FOUND
def test_estimate_scenario_returns_configured_projection(self, client: TestClient) -> None:
- """POST catalog estimate forwards request fields and returns the structured estimate."""
+ """Configured estimation returns the exact projection without touching run scheduling."""
estimate = ScenarioDefaultRunSizeEstimate(
status=ScenarioRunSizeEstimateStatus.Exact,
- total_attack_count=12,
- components=[ScenarioRunSizeComponent(label="Configured Jailbreak", count=12)],
+ total_attack_count=8,
+ components=[ScenarioRunSizeComponent(label="Configured Jailbreak", count=8)],
)
- with patch("pyrit.backend.routes.scenarios.get_scenario_service") as mock_get_service:
+ with (
+ patch("pyrit.backend.routes.scenarios.get_scenario_service") as mock_get_service,
+ patch("pyrit.backend.routes.scenarios.get_scenario_run_service") as mock_get_run_service,
+ ):
mock_service = MagicMock()
mock_service.estimate_scenario_run_size_async = AsyncMock(return_value=estimate)
mock_get_service.return_value = mock_service
@@ -780,7 +783,7 @@ def test_estimate_scenario_returns_configured_projection(self, client: TestClien
"/api/scenarios/catalog/airt.jailbreak/estimate",
json={
"techniques": ["prompt_sending"],
- "include_baseline": True,
+ "include_baseline": False,
"scenario_params": {
"num_jailbreaks": 2,
"num_jailbreak_attempts": 1,
@@ -789,14 +792,15 @@ def test_estimate_scenario_returns_configured_projection(self, client: TestClien
)
assert response.status_code == status.HTTP_200_OK
- assert response.json()["total_attack_count"] == 12
+ assert response.json()["total_attack_count"] == 8
request = mock_service.estimate_scenario_run_size_async.await_args.kwargs["request"]
assert request.techniques == ["prompt_sending"]
- assert request.include_baseline is True
+ assert request.include_baseline is False
assert request.scenario_params == {
"num_jailbreaks": 2,
"num_jailbreak_attempts": 1,
}
+ mock_get_run_service.assert_not_called()
async def test_estimate_scenario_supports_direct_keyword_call(self) -> None:
"""The FastAPI handler remains directly callable through its keyword-only API."""
diff --git a/tests/unit/exceptions/test_retry_collector.py b/tests/unit/exceptions/test_retry_collector.py
index f37b51a886..c4da265581 100644
--- a/tests/unit/exceptions/test_retry_collector.py
+++ b/tests/unit/exceptions/test_retry_collector.py
@@ -66,6 +66,25 @@ def test_record_extracts_exception_info(self) -> None:
assert evt.exception_type == "ValueError"
assert evt.exception_message == "test error"
+ def test_record_extracts_direct_or_response_status_code(self) -> None:
+ """record() preserves structured HTTP status codes without parsing messages."""
+ from unittest.mock import MagicMock
+
+ class DirectStatusError(Exception):
+ status_code = 429
+
+ class ResponseStatusError(Exception):
+ response = MagicMock(status_code=503)
+
+ collector = RetryCollector()
+ for exception in (DirectStatusError("limited"), ResponseStatusError("unavailable")):
+ retry_state = MagicMock(start_time=0.0, fn=None)
+ retry_state.outcome.failed = True
+ retry_state.outcome.exception.return_value = exception
+ collector.record(retry_state=retry_state)
+
+ assert [event.status_code for event in collector.events] == [429, 503]
+
def test_record_multiple_events(self) -> None:
"""record() accumulates events."""
from unittest.mock import MagicMock
diff --git a/tests/unit/memory/memory_interface/test_interface_scenario_history.py b/tests/unit/memory/memory_interface/test_interface_scenario_history.py
index 07c8d2de3d..c50523d962 100644
--- a/tests/unit/memory/memory_interface/test_interface_scenario_history.py
+++ b/tests/unit/memory/memory_interface/test_interface_scenario_history.py
@@ -5,14 +5,17 @@
import json
import uuid
+from contextlib import closing
from datetime import datetime, timedelta, timezone
-from unittest.mock import MagicMock
+from typing import Any
+from unittest.mock import MagicMock, patch
import pytest
+from sqlalchemy import and_, select, text
from unit.mocks import get_mock_target_identifier, make_scenario_result
-from pyrit.memory import MemoryInterface, ScenarioHistoryKeysetCursor
-from pyrit.memory.memory_models import ScenarioResultEntry
+from pyrit.memory import MemoryInterface, ScenarioHistoryKeysetCursor, SQLiteMemory
+from pyrit.memory.memory_models import AttackResultEntry, ScenarioResultEntry
from pyrit.models import (
SCENARIO_RUN_PLAN_METADATA_KEY,
AttackOutcome,
@@ -24,6 +27,23 @@
)
+class _LegacyScenarioLabelMemory(SQLiteMemory):
+ """Concrete backend retaining the pre-history single-value label hook."""
+
+ _get_scenario_result_labels_condition = MemoryInterface._get_scenario_result_labels_condition
+
+ def _get_scenario_result_label_condition(self, *, labels: dict[str, str]) -> Any:
+ conditions = []
+ for key, value in labels.items():
+ conditions.append(
+ text("json_extract(labels, :scenario_label_path_0) = :scenario_label_value_0").bindparams(
+ scenario_label_path_0=f'$."{key}"',
+ scenario_label_value_0=value,
+ )
+ )
+ return and_(*conditions)
+
+
@pytest.mark.parametrize(
("method_name", "kwargs"),
[
@@ -139,6 +159,11 @@ def test_history_filters_names_statuses_and_labels_without_hydration(
labels={"operator": "bob", "operation": "nightly", "team.name": "safety"},
)
sqlite_instance.add_scenario_results_to_memory(scenario_results=[included, excluded])
+ started_at = timestamp + timedelta(seconds=30)
+ sqlite_instance.update_scenario_metadata_fields(
+ scenario_result_id=str(included.id),
+ fields={"started_at": started_at.isoformat()},
+ )
attacks = [
AttackResult(
attack_result_id=str(uuid.UUID(int=12)),
@@ -155,6 +180,7 @@ def test_history_filters_names_statuses_and_labels_without_hydration(
},
error_type="RuntimeError",
error_message="failed",
+ total_retries=-3,
),
AttackResult(
attack_result_id=str(uuid.UUID(int=13)),
@@ -190,6 +216,7 @@ def test_history_filters_names_statuses_and_labels_without_hydration(
)
assert [row.scenario_result_id for row in rows] == [str(included.id)]
+ assert rows[0].started_at == started_at
assert rows[0].scenario_identifier["class_name"] == "ImplementationClass"
assert rows[0].scenario_registry_name == "registered.scenario"
compact_groups = (
@@ -217,6 +244,86 @@ def test_history_filters_names_statuses_and_labels_without_hydration(
assert has_more is False
+def test_legacy_label_hook_is_constructible_and_composes_multi_value_semantics(
+ sqlite_instance: MemoryInterface,
+) -> None:
+ timestamp = datetime(2026, 8, 7, tzinfo=timezone.utc)
+ included = _make_scenario(
+ result_id=uuid.UUID(int=30),
+ timestamp=timestamp,
+ name="Included",
+ state=ScenarioRunState.COMPLETED,
+ labels={"operator": "alice", "operation": "nightly"},
+ )
+ excluded = _make_scenario(
+ result_id=uuid.UUID(int=31),
+ timestamp=timestamp,
+ name="Excluded",
+ state=ScenarioRunState.COMPLETED,
+ labels={"operator": "carol", "operation": "nightly"},
+ )
+ sqlite_instance.add_scenario_results_to_memory(scenario_results=[included, excluded])
+ legacy = object.__new__(_LegacyScenarioLabelMemory)
+ condition = legacy._get_scenario_result_labels_condition(
+ labels={"operator": ["alice", "bob"], "operation": "nightly"}
+ )
+
+ with closing(sqlite_instance.get_session()) as session:
+ ids = session.execute(select(ScenarioResultEntry.id).where(condition)).scalars().all()
+
+ assert "_get_scenario_result_label_condition" not in _LegacyScenarioLabelMemory.__abstractmethods__
+ assert ids == [included.id]
+
+
+def test_nonterminal_state_projection_is_bounded_and_never_hydrates_results(
+ sqlite_instance: MemoryInterface,
+) -> None:
+ timestamp = datetime(2026, 8, 7, tzinfo=timezone.utc)
+ queued = _make_scenario(
+ result_id=uuid.UUID(int=40),
+ timestamp=timestamp,
+ name="Queued",
+ state=ScenarioRunState.QUEUED,
+ labels={},
+ )
+ running = _make_scenario(
+ result_id=uuid.UUID(int=41),
+ timestamp=timestamp,
+ name="Running",
+ state=ScenarioRunState.IN_PROGRESS,
+ labels={},
+ )
+ completed = _make_scenario(
+ result_id=uuid.UUID(int=42),
+ timestamp=timestamp,
+ name="Completed",
+ state=ScenarioRunState.COMPLETED,
+ labels={},
+ )
+ sqlite_instance.add_scenario_results_to_memory(scenario_results=[queued, running, completed])
+
+ with (
+ patch.object(ScenarioResultEntry, "get_scenario_result", side_effect=AssertionError("hydrated ScenarioResult")),
+ patch.object(AttackResultEntry, "get_attack_result", side_effect=AssertionError("hydrated AttackResult")),
+ ):
+ first, has_more = sqlite_instance.get_scenario_run_state_page(
+ states=[ScenarioRunState.QUEUED, ScenarioRunState.IN_PROGRESS],
+ limit=1,
+ )
+ second, second_has_more = sqlite_instance.get_scenario_run_state_page(
+ states=[ScenarioRunState.QUEUED, ScenarioRunState.IN_PROGRESS],
+ after_id=first[-1].scenario_result_id,
+ limit=1,
+ )
+
+ assert [record.state for record in [*first, *second]] == [
+ ScenarioRunState.QUEUED,
+ ScenarioRunState.IN_PROGRESS,
+ ]
+ assert has_more is True
+ assert second_has_more is False
+
+
def test_unique_scenario_labels_are_grouped_for_filter_options(sqlite_instance: MemoryInterface) -> None:
timestamp = datetime(2026, 8, 7, tzinfo=timezone.utc)
scenarios = [
diff --git a/tests/unit/memory/test_azure_sql_memory.py b/tests/unit/memory/test_azure_sql_memory.py
index 298a209fb2..dff3bbc779 100644
--- a/tests/unit/memory/test_azure_sql_memory.py
+++ b/tests/unit/memory/test_azure_sql_memory.py
@@ -442,7 +442,7 @@ def test_scenario_history_conditions_bind_or_within_label_and_registry_values(
memory_interface: AzureSQLMemory,
) -> None:
"""Scenario-history SQL Server conditions bind repeated values without interpolation."""
- label_condition = memory_interface._get_scenario_result_label_condition(
+ label_condition = memory_interface._get_scenario_result_labels_condition(
labels={"team.name": ["alice", "bob"], "operation": "nightly"}
)
registry_condition = memory_interface._get_scenario_registry_name_condition(
diff --git a/tests/unit/registry/test_scenario_registry.py b/tests/unit/registry/test_scenario_registry.py
index 393d6c5c17..b3abf2b2b0 100644
--- a/tests/unit/registry/test_scenario_registry.py
+++ b/tests/unit/registry/test_scenario_registry.py
@@ -122,6 +122,7 @@ async def test_create_and_initialize_async_creates_sets_params_and_initializes()
"my.scenario",
scenario_params={"foo": "bar"},
scenario_result_id="sr-1",
+ initial_metadata={"scheduler_managed_by": "test"},
objective_target=target,
max_concurrency=2,
)
@@ -129,6 +130,7 @@ async def test_create_and_initialize_async_creates_sets_params_and_initializes()
assert result is scenario
registry.create_instance.assert_called_once_with("my.scenario", scenario_result_id="sr-1")
scenario.set_scenario_registry_name.assert_called_once_with(scenario_registry_name="my.scenario")
+ scenario.set_initial_metadata.assert_called_once_with(metadata={"scheduler_managed_by": "test"})
scenario.set_params_from_args.assert_called_once_with(
args={"foo": "bar", "objective_target": target, "max_concurrency": 2}
)
diff --git a/tests/unit/scenario/core/test_scenario.py b/tests/unit/scenario/core/test_scenario.py
index d7f7fd7cd8..da2a405868 100644
--- a/tests/unit/scenario/core/test_scenario.py
+++ b/tests/unit/scenario/core/test_scenario.py
@@ -273,6 +273,7 @@ async def test_initialize_async_populates_atomic_attacks(self, mock_atomic_attac
assert scenario.atomic_attack_count == 0
scenario.set_params_from_args(args={"objective_target": mock_objective_target})
+ scenario.set_initial_metadata(metadata={"scheduler_managed_by": "test"})
await scenario.initialize_async()
assert scenario.atomic_attack_count == len(mock_atomic_attacks)
@@ -280,6 +281,7 @@ async def test_initialize_async_populates_atomic_attacks(self, mock_atomic_attac
[stored] = scenario._memory.get_scenario_results(scenario_result_ids=[scenario._scenario_result_id])
assert stored.metadata["run_plan"]["version"] == 1
assert len(stored.metadata["run_plan"]["atomic_groups"]) == len(mock_atomic_attacks)
+ assert stored.metadata["scheduler_managed_by"] == "test"
async def test_initialize_async_deduplicates_logical_seed_groups_in_run_plan(self, mock_objective_target) -> None:
duplicate_seed_groups = [
@@ -305,7 +307,7 @@ async def test_initialize_async_deduplicates_logical_seed_groups_in_run_plan(sel
expected_seed_id = duplicate_seed_groups[0].logical_id
assert persisted_plan["atomic_groups"][0]["seed_group_ids"] == [expected_seed_id]
assert [seed_group["id"] for seed_group in persisted_plan["seed_groups"]] == [expected_seed_id]
- assert scenario._build_run_plan().model_dump(mode="json") == persisted_plan
+ assert scenario._build_run_plan().model_dump(mode="json", exclude_none=True) == persisted_plan
assert atomic_attack.seed_groups is duplicate_seed_groups
assert len(atomic_attack.seed_groups) == 2
@@ -348,6 +350,24 @@ async def test_initialize_async_sets_objective_target(self, mock_objective_targe
assert scenario._objective_target_identifier.class_name == "MockTarget"
assert scenario._objective_target_identifier.class_module == "test"
+ async def test_initial_metadata_survives_subclass_metadata_override(self, mock_objective_target):
+ scenario = ConcreteScenario(name="Test Scenario", version=1)
+ scenario.set_params_from_args(args={"objective_target": mock_objective_target})
+ scenario.set_initial_metadata(metadata={"scheduler_managed_by": "test"})
+
+ with patch.object(
+ scenario,
+ "_build_initial_scenario_metadata",
+ return_value={"scenario_owned": "value"},
+ ):
+ await scenario.initialize_async()
+
+ [stored] = scenario._memory.get_scenario_results(scenario_result_ids=[scenario._scenario_result_id])
+ assert stored.metadata == {
+ "scenario_owned": "value",
+ "scheduler_managed_by": "test",
+ }
+
async def test_initialize_async_requires_objective_target(self):
"""Test that initialize_async raises ValueError when objective_target is None."""
scenario = ConcreteScenario(