+
+ 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
+
+
+ ) : (
@@ -452,8 +503,9 @@ function ScenarioRunPageContent({ scenarioResultId, attackResultId }: ScenarioRu
+ )}
- {isTerminalRunState(run.status) ? `Run ${formatRunState(run.status)}` : ''}
+ {queued ? progressText : isTerminalRunState(run.status) ? `Run ${formatRunState(run.status)}` : ''}
@@ -638,7 +690,9 @@ function ScenarioRunPageContent({ scenarioResultId, attackResultId }: ScenarioRu
Cancel this scenario run?
- In-flight work will be stopped. Finished executions 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 && (
@@ -650,6 +704,7 @@ function ScenarioRunPageContent({ scenarioResultId, attackResultId }: ScenarioRu
}
onClick={() => void handleCancel()}
@@ -819,3 +874,16 @@ function formatSuccess(succeeded: number, evaluated: number, percent: number | n
function formatCompletion(completed: number, planned: number | null): string {
return planned === null ? `${completed}/total unavailable` : `${completed}/${planned}`
}
+
+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())
+}
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 1758ab5b1b..b86484d01d 100644
--- a/frontend/src/services/api.ts
+++ b/frontend/src/services/api.ts
@@ -36,6 +36,7 @@ import type {
ScenarioRunSummary,
ScenarioRunListResponse,
ScenarioRunProgress,
+ ScenarioQueueSnapshot,
ScenarioRunState,
ConfigurationFileContent,
EnvironmentFileContent,
@@ -496,6 +497,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 9dab87b866..c8ab83d3d8 100644
--- a/frontend/src/types/index.ts
+++ b/frontend/src/types/index.ts
@@ -739,6 +739,7 @@ export interface RetryEvent {
component_role: string
component_name?: string | null
endpoint?: string | null
+ status_code?: number | null
elapsed_seconds: number
}
@@ -750,6 +751,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
@@ -757,6 +767,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
@@ -777,6 +788,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 {
@@ -793,6 +807,7 @@ export interface ScenarioRunListItem {
scenario_version: number
status: ScenarioRunState
created_at: string
+ started_at?: string | null
updated_at: string
error?: string | null
error_type?: string | null
@@ -826,6 +841,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
@@ -833,6 +849,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 2d1a526b0e..46f258ffd7 100644
--- a/frontend/src/utils/scenarioRunProgress.test.ts
+++ b/frontend/src/utils/scenarioRunProgress.test.ts
@@ -55,6 +55,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,
results: [],
@@ -88,6 +89,74 @@ describe('scenarioRunProgressReducer', () => {
expect(duplicate.summary).toEqual(updatedSummary)
})
+ it('treats cumulative overload snapshots as authoritative 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,
+ count: 2,
+ rate_limit_count: 2,
+ latest_timestamp: '2026-01-01T00:02:00Z',
+ }],
+ },
+ }),
+ fresh: false,
+ })
+ const repeated = scenarioRunProgressReducer(caughtUp, {
+ type: 'apply-page',
+ page: makePage({
+ plan: null,
+ run: caughtUp.run ?? makePage().run,
+ }),
+ 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)
+ expect(repeated.overloadSummaries[0].count).toBe(2)
+ })
+
it('retains last-good data and marks it stale after a transient failure', () => {
const first = scenarioRunProgressReducer(INITIAL_SCENARIO_RUN_PROGRESS_STATE, {
type: 'apply-page',
@@ -109,15 +178,39 @@ describe('scenarioRunProgressReducer', () => {
describe('scenario run timing', () => {
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)
+ 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-01T01:05:00Z'))).toBe(180_000)
+ expect(getEtaMilliseconds(active, SUMMARY.overall, 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,
}
- expect(getElapsedMilliseconds(terminal, Date.parse('2026-01-01T00:05:00Z'))).toBe(180_000)
+ const now = Date.parse('2026-01-01T01:00:00Z')
+
+ expect(getElapsedMilliseconds(queued, now)).toBe(0)
+
+ expect(getEtaMilliseconds(queued, SUMMARY.overall, now)).toBeNull()
+
+ expect(getElapsedMilliseconds(
+ { ...queued, status: 'CANCELLED', completed_at: '2026-01-01T00:30:00Z' },
+ now,
+ )).toBe(0)
})
it('calculates ETA from backend counts and hides unsafe estimates', () => {
diff --git a/frontend/src/utils/scenarioRunProgress.ts b/frontend/src/utils/scenarioRunProgress.ts
index 8edadc8f8d..1a701d14f3 100644
--- a/frontend/src/utils/scenarioRunProgress.ts
+++ b/frontend/src/utils/scenarioRunProgress.ts
@@ -2,6 +2,7 @@ import type {
ScenarioProgressCounts,
ScenarioProgressHeader,
ScenarioProgressResult,
+ ScenarioOverloadSummary,
ScenarioProgressSummary,
ScenarioRunPlan,
ScenarioRunState,
@@ -19,6 +20,7 @@ export interface ScenarioRunProgressState {
readonly results: ScenarioProgressResult[]
readonly error: string | null
readonly stale: boolean
+ readonly overloadSummaries: ScenarioOverloadSummary[]
}
export type ScenarioRunProgressAction =
@@ -38,6 +40,7 @@ export const INITIAL_SCENARIO_RUN_PROGRESS_STATE: ScenarioRunProgressState = {
results: [],
error: null,
stale: false,
+ overloadSummaries: [],
}
export function isTerminalRunState(status: ScenarioRunState): boolean {
@@ -78,6 +81,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,
@@ -85,9 +89,13 @@ 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 ?? [],
},
error: null,
stale: false,
+ overloadSummaries: state.overloadSummaries,
}
}
@@ -103,6 +111,9 @@ export function scenarioRunProgressReducer(
}
const results = [...resultsById.values()].sort(compareAttempts)
+ const overloadSummaries = [...(action.page.run.overload_summaries ?? [])].sort(
+ (left, right) => Date.parse(right.latest_timestamp) - Date.parse(left.latest_timestamp),
+ )
return {
loadStatus: 'ready',
run: action.page.run,
@@ -112,6 +123,7 @@ export function scenarioRunProgressReducer(
results,
error: null,
stale: false,
+ overloadSummaries,
}
}
@@ -119,15 +131,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/infra/README.md b/infra/README.md
index 82e09bdbe6..ab71992d91 100644
--- a/infra/README.md
+++ b/infra/README.md
@@ -666,7 +666,7 @@ Supported Azure integrations, including OpenAI, Content Safety, and Speech, can
- **Network outputs**: `egressPublicIpAddress`, `natGatewayId`, `acaInfrastructureSubnetId`, and `vnetName` describe the created network.
- **PIP lock**: `protectEgressPublicIp=true` creates a resource-scoped `CanNotDelete` lock. The internal ADO workflow enables it; community examples leave it disabled unless the operator explicitly opts in.
- **Log Analytics shared key**: `listKeys()` is the standard ACA pattern. The key is used during deployment only, not exposed to the application.
-- **Workload profiles**: Consumption tier. Defaults to 1 replica (no auto-scale).
+- **Workload profiles**: Consumption tier. Scenario FIFO admission is process-local, so deployments allow at most 1 active replica and default to 1 (no auto-scale). Supporting multiple replicas requires shared, database-backed admission or lease ownership before raising this limit.
- **Key Vault**: Bicep requires a supplied vault resource ID. The vault is backup/audit-only for `deploy_instance.py`, but it is the editable runtime source for deployments using `envSecretName`. Those app identities require `Key Vault Secrets Officer` and a permitted network path. AcrPull is still granted separately.
- **OpenTelemetry**: When `enableOtel=true`, configure the agent post-deploy:
```bash
diff --git a/infra/application.bicep b/infra/application.bicep
index 974a73b221..2784153406 100644
--- a/infra/application.bicep
+++ b/infra/application.bicep
@@ -56,7 +56,10 @@ param memoryGb string = '2.0'
@description('Minimum number of replicas')
param minReplicas int = 1
-@description('Maximum number of replicas')
+@description('Maximum number of replicas. Must remain 1 while scenario FIFO scheduling is process-local.')
+@allowed([
+ 1
+])
param maxReplicas int = 1
@description('CIDR range allowed to reach ACA directly. Empty = unrestricted. Must be empty when Front Door is enabled because ACA sees Front Door backend IPs, not client IPs.')
diff --git a/pyrit/backend/main.py b/pyrit/backend/main.py
index 39d38dee17..70afaefea4 100644
--- a/pyrit/backend/main.py
+++ b/pyrit/backend/main.py
@@ -41,6 +41,7 @@
from pyrit.backend.services.configuration_file_service import ConfigurationFileService
from pyrit.backend.services.converter_service import get_converter_service
from pyrit.backend.services.environment_file_service import EnvironmentFileService
+from pyrit.backend.services.scenario_run_service import get_scenario_run_service
from pyrit.common.path import CONFIGURATION_DIRECTORY_PATH
from pyrit.registry import InitializerRegistry
from pyrit.setup.configuration_loader import ConfigurationLoader
@@ -106,6 +107,9 @@ 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.
@@ -116,9 +120,12 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
yield
finally:
try:
- await converter_service.close_async()
+ await scenario_run_service.shutdown_async()
finally:
- get_converter_service.cache_clear()
+ try:
+ await converter_service.close_async()
+ finally:
+ get_converter_service.cache_clear()
app = FastAPI(
diff --git a/pyrit/backend/routes/scenarios.py b/pyrit/backend/routes/scenarios.py
index 7905a7f72b..bb97064a5f 100644
--- a/pyrit/backend/routes/scenarios.py
+++ b/pyrit/backend/routes/scenarios.py
@@ -23,15 +23,14 @@
from pyrit.backend.routes.common import parse_label_query_params
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,
ScenarioRunSizeEstimate,
ScenarioRunSizeEstimateRequest,
ScenarioRunSummary,
)
-from pyrit.models.scenario_progress import ScenarioRunProgress
router = APIRouter(prefix="/scenarios", tags=["scenarios"])
@@ -214,6 +213,20 @@ async def list_scenario_runs( # pyrit-async-suffix-exempt
)
+@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,
@@ -237,6 +250,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(
@@ -275,6 +290,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 1a8098d69a..fb17423d3b 100644
--- a/pyrit/backend/services/scenario_run_service.py
+++ b/pyrit/backend/services/scenario_run_service.py
@@ -15,10 +15,11 @@
import json
import logging
import uuid
+from collections import OrderedDict, deque
from collections.abc import Mapping, Sequence
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
-from datetime import datetime
+from datetime import UTC, datetime
from threading import Lock
from typing import Any
from urllib.parse import urlsplit, urlunsplit
@@ -35,7 +36,7 @@
)
from pyrit.backend.services.scenario_configuration_resolver import ScenarioConfigurationResolver
from pyrit.backend.services.scenario_progress_read_model import ResultUnitIdentity, ScenarioProgressReadModel
-from pyrit.memory import AttackResultKeysetCursor, CentralMemory
+from pyrit.memory import AttackResultKeysetCursor, CentralMemory, SQLiteMemory
from pyrit.memory.memory_interface import (
ScenarioHistoryAggregate,
ScenarioHistoryKeysetCursor,
@@ -43,11 +44,14 @@
)
from pyrit.models import (
SCENARIO_RUN_PLAN_METADATA_KEY,
+ SCENARIO_RUN_STARTED_AT_METADATA_KEY,
AttackOutcome,
ComponentIdentifier,
ScenarioAttackResultDelta,
ScenarioIdentifier,
ScenarioProgressHeader,
+ ScenarioQueueEntry,
+ ScenarioQueueSnapshot,
ScenarioResult,
ScenarioRunPlan,
ScenarioRunPlanAtomicGroup,
@@ -59,6 +63,7 @@
AttackErrorSummary,
AttackRetrySummary,
RunScenarioRequest,
+ ScenarioOverloadSummary,
ScenarioRunListItem,
ScenarioRunSummary,
ScenarioTargetSummary,
@@ -69,7 +74,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"
_SAFE_SCENARIO_PARAMETER_NAMES = frozenset(
{
@@ -85,6 +104,7 @@
)
_HISTORY_ATOMIC_GROUPS_ADAPTER = TypeAdapter(list[ScenarioRunPlanAtomicGroup])
_HISTORY_SEED_ID_MAP_ADAPTER = TypeAdapter(list[dict[str, str]])
+_STARTED_AT_ADAPTER = TypeAdapter(datetime)
@dataclass
@@ -95,6 +115,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)
@@ -103,6 +132,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:
@@ -110,7 +141,10 @@ 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.
+ FIFO ordering therefore spans only runs submitted to the same backend
+ process. Deploy one backend replica to preserve a global admission order;
+ multiple replicas require a shared database-backed scheduler or lease.
"""
#: Seconds to let initialization's own background tasks (for example HTTP client teardown
@@ -119,11 +153,16 @@ class ScenarioRunService:
_INITIALIZATION_DRAIN_TIMEOUT = 5.0
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._configuration_resolver = ScenarioConfigurationResolver()
self._progress_read_model = ScenarioProgressReadModel(memory=self._memory)
self._technique_metadata_cache: dict[str, dict[str, ScenarioTechniqueSummary]] = {}
@@ -135,115 +174,106 @@ def __init__(self, *, max_concurrent_runs: int = _DEFAULT_MAX_CONCURRENT_RUNS) -
# they are serialized onto a single worker. The event loop is still free while they run,
# which is the point of the offload.
self._prepare_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="pyrit-scenario-prep")
+ 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()
-
- # This frame owns the permit until the background task is created; every exit path
- # before that hand-off has to release it, including cancellation, which is a
- # BaseException and so is not caught by ``except Exception``.
- release_on_exit = True
- registered_run_id: str | None = None
- try:
- # A resumed run keeps the state its previous run left behind, so one that was
- # cancelled and is now being resumed on purpose is still CANCELLED while it
- # initializes. Read that before preparation: the check afterwards otherwise
- # cannot tell an intentional resume from a cancellation that landed while the
- # worker thread was still initializing, and would refuse to restart it.
+ async with self._launch_lock:
+ if self._stopping:
+ raise RuntimeError("Scenario run scheduling is stopping.")
resumed_from_cancelled = self._is_run_cancelled(scenario_result_id=request.scenario_result_id)
-
- # Initialization loads the default datasets, which takes minutes, and is mostly
- # synchronous work. Run it on a worker thread so the event loop stays free to
- # answer health checks and status polls while a run is starting.
prepare_task = asyncio.get_running_loop().run_in_executor(
- self._prepare_executor, functools.partial(self._prepare_run_blocking, request=request)
+ self._prepare_executor,
+ functools.partial(self._prepare_run_blocking, request=request),
)
try:
scenario = await asyncio.shield(prepare_task)
- except BaseException as exc:
- # A worker thread cannot be killed, so it keeps initializing after this frame
- # unwinds. Keep holding the permit until it actually finishes, otherwise the
- # next caller is admitted while this run is still loading datasets and
- # ``max_concurrent_runs`` stops bounding the work that is really running.
- if not prepare_task.done():
- prepare_task.add_done_callback(self._release_abandoned_prepare)
- release_on_exit = False
- elif isinstance(exc, asyncio.CancelledError):
- # The thread can finish just as the cancellation lands. A done future never
- # calls back, so cleaning up here is the only chance to release the permit
- # and terminalize the run that initialization already stored.
- release_on_exit = False
+ except asyncio.CancelledError:
+ if prepare_task.done():
try:
self._release_abandoned_prepare(prepare_task)
except Exception as cleanup_error:
- # The permit is released first, so it is already back even if the rest
- # failed. Never let cleanup replace the cancellation being propagated.
logger.warning(f"Could not clean up after a cancelled scenario preparation: {cleanup_error}")
+ else:
+ prepare_task.add_done_callback(self._release_abandoned_prepare)
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
- registered_run_id = scenario_result_id
-
- # Build the response before spawning the task so that a failure here cannot leave
- # a run executing that the caller never received an id for.
- response = self.get_run(scenario_result_id=scenario_result_id)
- if response is None:
+ 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.")
+ if not resumed_from_cancelled and persisted[0].scenario_run_state == ScenarioRunState.CANCELLED:
+ response = self._build_response(
+ scenario_result_id=scenario_result_id,
+ active_error=None,
+ queue_position=None,
+ active_scenario_result_id=self._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
+ if (
+ self._build_response(
+ scenario_result_id=scenario_result_id,
+ active_error=None,
+ queue_position=None,
+ active_scenario_result_id=self._active_scenario_result_id,
+ )
+ is None
+ ):
raise RuntimeError(
f"Scenario run {scenario_result_id} was not found in the database after 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(UTC),
+ )
+ await self._enqueue_run_async(scheduled=scheduled)
- # A run can be cancelled through its id while initialization is still on the worker
- # thread: a resume already knows the id, and a fresh run appears in the run list as
- # soon as initialization stores it. Nothing has run yet, so honour that instead of
- # starting a scenario the caller gave up on. The finally block returns the permit
- # and drops the tracking entry.
- if response.status == ScenarioRunState.CANCELLED and not resumed_from_cancelled:
- logger.info(f"Scenario run {scenario_result_id} was cancelled while it was being initialized.")
- return response
-
- # Spawn background task (only runs scenario.run_async). It releases the permit in
- # its own finally, so ownership transfers here and this frame must not release it.
- task = asyncio.create_task(self._execute_run_async(scenario_result_id=scenario_result_id))
- active.task = task
- release_on_exit = False
- registered_run_id = None
- finally:
- if registered_run_id is not None:
- self._active_tasks.pop(registered_run_id, None)
- if release_on_exit:
- self._run_semaphore.release()
-
+ 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
def _is_run_cancelled(self, *, scenario_result_id: str | None) -> bool:
@@ -268,17 +298,13 @@ def _release_abandoned_prepare(self, prepare_task: "asyncio.Future[Scenario]") -
"""
Clean up after an abandoned preparation thread has finished.
- ``start_run_async`` hands ownership of the permit to this callback when it is
- cancelled while the worker thread is still initializing, so the permit is only
- released after the thread has genuinely stopped using the slot. A preparation that
- succeeds anyway leaves behind a scenario result nobody will run, which is marked
- cancelled here rather than left waiting in ``CREATED``.
+ A preparation that succeeds after its caller is cancelled leaves behind a
+ scenario result nobody will run. Mark it cancelled rather than leaving it
+ in ``CREATED``.
Args:
prepare_task: The future wrapping the abandoned ``_prepare_run_blocking`` call.
"""
- self._run_semaphore.release()
-
if prepare_task.cancelled():
return
error = prepare_task.exception()
@@ -428,13 +454,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.
@@ -442,11 +475,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,
@@ -556,35 +596,400 @@ 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, TimeoutError):
- await asyncio.wait_for(active.task, timeout=5.0)
-
- # The run can reach a terminal state during the await above, so only cancel a run that
- # is still going. The re-read below reports whichever state actually won.
- self._memory.try_update_scenario_run_state(
+ 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.try_update_scenario_run_state,
+ scenario_result_id=scenario_result_id,
+ expected_states={ScenarioRunState.CREATED, ScenarioRunState.QUEUED},
+ 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.try_update_scenario_run_state,
+ scenario_result_id=scenario_result_id,
+ expected_states={
+ ScenarioRunState.CREATED,
+ ScenarioRunState.QUEUED,
+ ScenarioRunState.IN_PROGRESS,
+ },
+ scenario_run_state=ScenarioRunState.CANCELLED,
+ error_message=_USER_CANCELLATION_REASON,
+ error_type="CancelledError",
+ )
+
+ 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,
- expected_states={ScenarioRunState.CREATED, ScenarioRunState.IN_PROGRESS},
- 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(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,
)
- return self.get_run(scenario_result_id=scenario_result_id)
+ 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._launch_lock:
+ 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
+ await asyncio.to_thread(self._prepare_executor.shutdown, wait=True)
+ 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(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)
+
+ def _can_retry_active_run(self, *, scenario_result_id: str) -> bool:
+ """Return whether retry work may continue for the active run."""
+ return not self._stopping and self._active_scenario_result_id == scenario_result_id
+
+ 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 self._can_retry_active_run(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 self._can_retry_active_run(scenario_result_id=active.scenario_result_id):
+ await asyncio.sleep(delay)
+ try:
+ async with self._scheduler_lock:
+ if not self._can_retry_active_run(scenario_result_id=active.scenario_result_id):
+ return
+ await asyncio.to_thread(
+ self._memory.try_update_scenario_run_state,
+ scenario_result_id=active.scenario_result_id,
+ expected_states={ScenarioRunState.CREATED, ScenarioRunState.IN_PROGRESS},
+ 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,
+ )
+
+ @staticmethod
+ def _is_terminal_state(state: ScenarioRunState) -> bool:
+ """Return whether a scenario state is terminal."""
+ return state in (ScenarioRunState.COMPLETED, ScenarioRunState.FAILED, ScenarioRunState.CANCELLED)
async def _run_initializers_async(self, *, request: RunScenarioRequest) -> None:
"""
@@ -633,6 +1038,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,
)
@@ -642,36 +1048,70 @@ 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.try_update_scenario_run_state,
+ scenario_result_id=scenario_result_id,
+ expected_states={ScenarioRunState.CREATED, ScenarioRunState.IN_PROGRESS},
+ 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.try_update_scenario_run_state,
+ scenario_result_id=scenario_result_id,
+ expected_states={ScenarioRunState.CREATED, ScenarioRunState.IN_PROGRESS},
+ 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.
@@ -679,6 +1119,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.
@@ -686,13 +1128,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.
@@ -700,6 +1149,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.
@@ -763,6 +1214,7 @@ def _build_response_from_db(
failed_attacks: list[AttackErrorSummary] = []
attack_retries: list[AttackRetrySummary] = []
persisted_retries: list[int] = []
+ overload_events: deque[Any] = deque(maxlen=_MAX_OVERLOAD_EVENTS)
attempts_by_unit: dict[ResultUnitIdentity, int] = {}
for atomic_attack_name, results in scenario_result.attack_results.items():
for attack_result in results:
@@ -778,6 +1230,7 @@ def _build_response_from_db(
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),
@@ -793,7 +1246,7 @@ 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 = self._progress_read_model.total_retry_pressure(
@@ -812,6 +1265,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,
@@ -835,6 +1289,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),
)
@staticmethod
@@ -956,6 +1413,7 @@ def _build_history_summary(
scenario_version=record.scenario_version,
status=status,
created_at=record.created_at,
+ started_at=record.started_at,
updated_at=max(timestamps),
error=record.error_message,
error_type=record.error_type,
@@ -976,6 +1434,23 @@ def _build_history_summary(
attack_details_available=False,
)
+ @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 _identifier_techniques(scenario_identifier: ScenarioIdentifier | None) -> list[str]:
"""
@@ -1008,6 +1483,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:
"""
@@ -1071,10 +1595,16 @@ def _safe_endpoint(endpoint: str | None) -> str | None:
return urlunsplit((parsed.scheme, host, "", "", ""))
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:
@@ -1084,11 +1614,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,
+ )
def _load_run_plan(self, *, scenario_result: ScenarioResult) -> ScenarioRunPlan | None:
"""
@@ -1192,6 +1740,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(
@@ -1201,6 +1751,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)
@@ -1233,6 +1785,9 @@ def get_run_progress_from_storage(
terminal=terminal,
objective_scorer_identifier=objective_scorer_identifier,
)
+ overload_events: deque[Any] = deque(maxlen=_MAX_OVERLOAD_EVENTS)
+ for delta in progress_snapshot.deltas:
+ overload_events.extend(delta.retry_events)
available = [
(delta, result)
for delta, result in zip(progress_snapshot.deltas, progress_snapshot.results, strict=True)
@@ -1262,6 +1817,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,
@@ -1269,6 +1825,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,
results=results,
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 6c53acb553..42aae06309 100644
--- a/pyrit/memory/__init__.py
+++ b/pyrit/memory/__init__.py
@@ -22,6 +22,7 @@
ScenarioHistoryAggregate,
ScenarioHistoryKeysetCursor,
ScenarioHistoryRunRecord,
+ ScenarioRunStateRecord,
)
from pyrit.memory.memory_models import AttackResultEntry, EmbeddingDataEntry, PromptMemoryEntry, SeedEntry
from pyrit.memory.sqlite_memory import SQLiteMemory
@@ -64,6 +65,7 @@
"ScenarioHistoryKeysetCursor": "pyrit.memory.memory_interface",
"ScenarioHistoryRunRecord": "pyrit.memory.memory_interface",
"ScenarioHistoryAggregate": "pyrit.memory.memory_interface",
+ "ScenarioRunStateRecord": "pyrit.memory.memory_interface",
"PromptMemoryEntry": "pyrit.memory.memory_models",
"SeedEntry": "pyrit.memory.memory_models",
"set_message_piece_sha256_async": "pyrit.memory.storage",
diff --git a/pyrit/memory/azure_sql_memory.py b/pyrit/memory/azure_sql_memory.py
index 1c7dd0cb2d..651bfb76aa 100644
--- a/pyrit/memory/azure_sql_memory.py
+++ b/pyrit/memory/azure_sql_memory.py
@@ -611,7 +611,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.
@@ -708,6 +726,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 949b31a341..12808ed512 100644
--- a/pyrit/memory/memory_interface.py
+++ b/pyrit/memory/memory_interface.py
@@ -202,6 +202,7 @@ 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)
@@ -244,6 +245,14 @@ def empty(cls, *, scenario_result_id: str) -> "ScenarioHistoryAggregate":
)
+@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)
class _AttackResultQuery:
"""
@@ -1727,17 +1736,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.
@@ -1762,6 +1792,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.
@@ -4588,6 +4622,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.
"""
@@ -4600,6 +4657,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,
@@ -4699,12 +4759,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,
*,
@@ -4756,7 +4876,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(
@@ -4779,6 +4899,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,
@@ -4812,6 +4933,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,
@@ -4922,7 +5044,13 @@ def _build_scenario_history_aggregate_statement(
AttackResultEntry.objective_sha256.label("objective_sha256"),
AttackResultEntry.outcome.label("outcome"),
AttackResultEntry.timestamp.label("timestamp"),
- func.coalesce(AttackResultEntry.total_retries, 0).label("total_retries"),
+ case(
+ (
+ func.coalesce(AttackResultEntry.total_retries, 0) > 0,
+ func.coalesce(AttackResultEntry.total_retries, 0),
+ ),
+ else_=0,
+ ).label("total_retries"),
)
.where(AttackResultEntry.attribution_parent_id.in_(entry_ids))
.subquery("history_attempts")
@@ -5074,6 +5202,22 @@ def _build_scenario_history_unit_statement(self, *, attempts: Any, plan_entry_id
).label("is_planned"),
).where(matched.c.match_rank == 1)
+ @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 9d8066a745..04421c09f3 100644
--- a/pyrit/memory/sqlite_memory.py
+++ b/pyrit/memory/sqlite_memory.py
@@ -484,9 +484,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:
@@ -558,6 +569,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 0596833912..7ae1bc7c78 100644
--- a/pyrit/models/__init__.py
+++ b/pyrit/models/__init__.py
@@ -108,6 +108,7 @@
from pyrit.models.scenario_progress import (
SCENARIO_RUN_PLAN_METADATA_KEY,
SCENARIO_RUN_PLAN_VERSION,
+ SCENARIO_RUN_STARTED_AT_METADATA_KEY,
ScenarioAtomicGroupProgress,
ScenarioAttackResultDelta,
ScenarioAttackTechniqueDetails,
@@ -120,6 +121,8 @@
ScenarioProgressResult,
ScenarioProgressScore,
ScenarioProgressSummary,
+ ScenarioQueueEntry,
+ ScenarioQueueSnapshot,
ScenarioRunPlan,
ScenarioRunPlanAtomicGroup,
ScenarioRunPlanSeedGroup,
@@ -279,6 +282,7 @@
"ScenarioRunState": "pyrit.models.results.scenario_result",
"SCENARIO_RUN_PLAN_METADATA_KEY": "pyrit.models.scenario_progress",
"SCENARIO_RUN_PLAN_VERSION": "pyrit.models.scenario_progress",
+ "SCENARIO_RUN_STARTED_AT_METADATA_KEY": "pyrit.models.scenario_progress",
"ScenarioAttackResultDelta": "pyrit.models.scenario_progress",
"ScenarioAtomicGroupProgress": "pyrit.models.scenario_progress",
"ScenarioAttackTechniqueDetails": "pyrit.models.scenario_progress",
@@ -291,6 +295,8 @@
"ScenarioProgressResult": "pyrit.models.scenario_progress",
"ScenarioProgressScore": "pyrit.models.scenario_progress",
"ScenarioProgressSummary": "pyrit.models.scenario_progress",
+ "ScenarioQueueEntry": "pyrit.models.scenario_progress",
+ "ScenarioQueueSnapshot": "pyrit.models.scenario_progress",
"ScenarioRunPlan": "pyrit.models.scenario_progress",
"ScenarioRunPlanAtomicGroup": "pyrit.models.scenario_progress",
"ScenarioRunPlanSeedPrompt": "pyrit.models.scenario_progress",
diff --git a/pyrit/models/catalog/scenario.py b/pyrit/models/catalog/scenario.py
index 22c017a761..665dcd26f7 100644
--- a/pyrit/models/catalog/scenario.py
+++ b/pyrit/models/catalog/scenario.py
@@ -310,6 +310,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)."""
@@ -319,6 +330,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")
@@ -337,7 +349,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")
@@ -358,6 +373,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 ScenarioRunListItem(BaseModel):
@@ -369,6 +390,7 @@ class ScenarioRunListItem(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="Persisted run-level error message")
error_type: str | None = Field(None, description="Persisted run-level exception class")
diff --git a/pyrit/models/results/scenario_result.py b/pyrit/models/results/scenario_result.py
index 9c4f829141..e13ca5613e 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 64e6470d96..cdbe1bbefa 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 0122f14e3f..33a9f48fa0 100644
--- a/pyrit/models/scenario_progress.py
+++ b/pyrit/models/scenario_progress.py
@@ -8,7 +8,7 @@
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
@@ -16,6 +16,7 @@
from pyrit.models.score.score import ScoreStatus
SCENARIO_RUN_PLAN_METADATA_KEY = "run_plan"
+SCENARIO_RUN_STARTED_AT_METADATA_KEY = "started_at"
SCENARIO_RUN_PLAN_VERSION = 1
@@ -100,6 +101,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
@@ -107,6 +109,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 ScenarioProgressScore(BaseModel):
@@ -246,6 +251,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 e5823ec1f3..cbd84f5206 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
@@ -245,6 +246,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:
"""
@@ -274,6 +276,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``).
@@ -287,5 +291,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 8e87d288ca..b61440dd11 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, final
@@ -229,6 +229,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
@@ -276,6 +277,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]:
"""
@@ -979,7 +984,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 6a8f883a5e..d9790b81e7 100644
--- a/tests/unit/backend/test_main.py
+++ b/tests/unit/backend/test_main.py
@@ -20,16 +20,32 @@
from pyrit.backend.main import SPAStaticFiles, app, lifespan, setup_frontend
from pyrit.backend.models.converters import CreateConverterRequest
-from pyrit.backend.services.converter_service import get_converter_service
+from pyrit.backend.services.converter_service import ConverterService, get_converter_service
+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."""
@pytest.mark.parametrize("fail_during_lifespan", [False, True])
- async def test_lifespan_cleans_converter_uploads(self, fail_during_lifespan: bool) -> None:
+ async def test_lifespan_cleans_converter_uploads(
+ self, fail_during_lifespan: bool, mock_scenario_run_lifecycle
+ ) -> None:
fake_config = ConfigurationLoader()
+ created_resources: tuple[ConverterService, Path] | None = None
with (
patch.object(ConfigurationLoader, "load_with_overrides", return_value=fake_config),
patch.object(ConfigurationLoader, "initialize_pyrit_async", new=AsyncMock()),
@@ -48,16 +64,19 @@ async def test_lifespan_cleans_converter_uploads(self, fail_during_lifespan: boo
entry = service._registry.instances.get_entry("lifespan-upload")
assert entry is not None
owned_path = Path(entry.metadata["owned_artifact_paths"][0])
+ created_resources = (service, owned_path)
assert owned_path.read_bytes() == b"%PDF-1.4\n"
if fail_during_lifespan:
raise RuntimeError("application failed")
+ assert created_resources is not None
+ service, owned_path = created_resources
assert not owned_path.exists()
assert not service._upload_path.exists()
assert service._registry.instances.get("lifespan-upload") is None
assert get_converter_service.cache_info().currsize == 0
- async def test_lifespan_restarts_with_fresh_upload_directory(self) -> None:
+ async def test_lifespan_restarts_with_fresh_upload_directory(self, mock_scenario_run_lifecycle) -> None:
fake_config = ConfigurationLoader()
paths: list[Path] = []
with (
@@ -71,10 +90,9 @@ async def test_lifespan_restarts_with_fresh_upload_directory(self) -> None:
assert path.is_dir()
paths.append(path)
assert not path.exists()
-
assert paths[0] != paths[1]
- 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 (
@@ -89,8 +107,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 (
@@ -104,7 +124,30 @@ 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_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 (
@@ -117,7 +160,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_exposes_configured_initializers(self) -> None:
+ async def test_lifespan_exposes_configured_initializers(self, mock_scenario_run_lifecycle) -> None:
"""Test that the active config initializer sequence is exposed to API routes."""
fake_config = ConfigurationLoader(
initializers=[
@@ -137,7 +180,7 @@ async def test_lifespan_exposes_configured_initializers(self) -> None:
assert app.state.configured_initializers[0].parameters == {"tags": ["default"]}
assert [item.order_index for item in app.state.configured_initializers] == [0, 1]
- async def test_lifespan_loads_explicit_config_as_override(self) -> None:
+ async def test_lifespan_loads_explicit_config_as_override(self, mock_scenario_run_lifecycle) -> None:
"""Test that PYRIT_CONFIG_FILE overlays the default configuration."""
fake_config = ConfigurationLoader()
with (
@@ -151,7 +194,7 @@ async def test_lifespan_loads_explicit_config_as_override(self) -> None:
assert str(load_mock.call_args.kwargs["config_file"]).endswith("foo.yaml")
- async def test_lifespan_configures_custom_initializer_source_from_config(self) -> None:
+ async def test_lifespan_configures_custom_initializer_source_from_config(self, mock_scenario_run_lifecycle) -> None:
"""Test that YAML config determines the custom script source."""
fake_config = ConfigurationLoader(custom_initializers_source="C:/yaml/initializers")
registry = MagicMock()
@@ -167,7 +210,7 @@ async def test_lifespan_configures_custom_initializer_source_from_config(self) -
registry.configure_custom_scripts_source.assert_called_once_with("C:/yaml/initializers")
registry.register_stored_initializers.assert_not_called()
- async def test_lifespan_registers_stored_initializers_when_enabled(self) -> None:
+ async def test_lifespan_registers_stored_initializers_when_enabled(self, mock_scenario_run_lifecycle) -> None:
"""Test that enabled custom initializers are registered before configured initialization."""
fake_config = ConfigurationLoader(allow_custom_initializers=True)
call_order: list[str] = []
@@ -190,7 +233,7 @@ async def initialize_async(*, raise_on_initializer_error: bool) -> None:
registry.register_stored_initializers.assert_called_once_with()
assert call_order == ["custom", "configured"]
- async def test_lifespan_downloads_blob_config_to_temporary_file(self) -> None:
+ async def test_lifespan_downloads_blob_config_to_temporary_file(self, mock_scenario_run_lifecycle) -> None:
"""Test that an Azure Blob config URI is materialized and removed after loading."""
fake_config = ConfigurationLoader()
config_content = b"operator: blob-user\n"
diff --git a/tests/unit/backend/test_scenario_run_routes.py b/tests/unit/backend/test_scenario_run_routes.py
index aac24e33e2..9f31b6f9b4 100644
--- a/tests/unit/backend/test_scenario_run_routes.py
+++ b/tests/unit/backend/test_scenario_run_routes.py
@@ -25,6 +25,8 @@
ScenarioProgressCounts,
ScenarioProgressHeader,
ScenarioProgressSummary,
+ ScenarioQueueEntry,
+ ScenarioQueueSnapshot,
ScenarioRunPlan,
ScenarioRunProgress,
ScenarioRunState,
@@ -199,7 +201,7 @@ def test_list_runs_rejects_unbounded_limit(self, client: TestClient) -> None:
async def test_list_runs_requires_keyword_arguments(self) -> None:
"""Test that route parameters cannot be passed positionally."""
with pytest.raises(TypeError, match="positional"):
- await list_scenario_runs(None, None, None, 100, None)
+ await list_scenario_runs(None, None, None, 100, None) # ty: ignore[too-many-positional-arguments]
def test_list_runs_returns_multiple_runs(self, client: TestClient) -> None:
"""Test that list runs returns all tracked runs."""
@@ -251,6 +253,45 @@ def test_list_runs_passes_repeated_filters_and_labels(self, client: TestClient)
)
+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=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}."""
@@ -358,7 +399,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
@@ -374,6 +420,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]
@@ -405,7 +453,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
@@ -421,6 +473,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 1ba888f798..623367c5db 100644
--- a/tests/unit/backend/test_scenario_run_service.py
+++ b/tests/unit/backend/test_scenario_run_service.py
@@ -11,7 +11,7 @@
import time
import uuid
from dataclasses import replace
-from datetime import UTC, datetime
+from datetime import UTC, datetime, timedelta
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
@@ -22,12 +22,18 @@
import pyrit.backend.services.scenario_run_service as _svc_mod
from pyrit.backend.services.scenario_progress_read_model import ScenarioPlanLookup, ScenarioProgressReadModel
from pyrit.backend.services.scenario_run_service import (
- _DEFAULT_MAX_CONCURRENT_RUNS,
ScenarioRunService,
)
from pyrit.common.utils import to_sha256
from pyrit.converter import Converter
-from pyrit.memory import AttackResultKeysetCursor, ScenarioHistoryAggregate, ScenarioHistoryRunRecord
+from pyrit.memory import (
+ AttackResultKeysetCursor,
+ AzureSQLMemory,
+ ScenarioHistoryAggregate,
+ ScenarioHistoryRunRecord,
+ ScenarioRunStateRecord,
+ SQLiteMemory,
+)
from pyrit.models import (
SCENARIO_RUN_PLAN_METADATA_KEY,
AtomicAttackIdentifier,
@@ -36,6 +42,7 @@
AttackSeedGroup,
AttackTechniqueIdentifier,
ComponentIdentifier,
+ RetryEvent,
ScenarioAttackResultDelta,
ScenarioProgressResult,
ScenarioProgressScore,
@@ -161,6 +168,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,
@@ -173,7 +181,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 = []
mock.get_scenario_run_history_page.return_value = ([], {}, False)
mock.get_scenario_history_aggregates.return_value = {}
@@ -231,15 +239,44 @@ def mock_all_registries(mock_memory):
class TestScenarioRunServiceStartRun:
"""Tests for ScenarioRunService.start_run_async."""
+ def test_init_rejects_nonpositive_max_concurrent_runs(self, mock_memory) -> None:
+ with pytest.raises(ValueError, match="at least 1"):
+ ScenarioRunService(max_concurrent_runs=0)
+
+ async def test_start_run_rejects_after_shutdown(self, mock_all_registries) -> None:
+ service = ScenarioRunService()
+ await service.shutdown_async()
+
+ with pytest.raises(RuntimeError, match="scheduling is stopping"):
+ await service.start_run_async(request=_make_request())
+
+ mock_all_registries["scenario_registry"].create_and_initialize_async.assert_not_awaited()
+
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."""
@@ -388,6 +425,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,
@@ -395,6 +433,113 @@ 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()
+ release_queued = 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()
+ await release_queued.wait()
+ 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.try_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)
+ queued_task = service._active_tasks["run-2"].task
+ assert queued_task is not None
+ release_queued.set()
+ await asyncio.wait_for(queued_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()
@@ -583,43 +728,113 @@ 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.try_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),
+ ]
- # A real run holds its permit until it finishes, so the background task has to stay
- # in flight for the limit to be reachable. The default AsyncMock returns immediately
- # and would hand every permit straight back.
- still_running = asyncio.Event()
-
- async def _block_until_released() -> None:
- await still_running.wait()
-
- scenario_instance.run_async = _block_until_released
-
- # 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)
+ ]
- try:
- # 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())
- finally:
- still_running.set()
+ 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."""
@@ -760,7 +975,6 @@ async def _complete_then_cancel(awaitable):
ScenarioRunState.CREATED,
ScenarioRunState.IN_PROGRESS,
}
- assert service._run_semaphore._value == _DEFAULT_MAX_CONCURRENT_RUNS
async def test_start_run_does_not_run_a_scenario_cancelled_during_initialization(self, mock_all_registries) -> None:
"""A run appears in the run list as soon as it is stored, so it can be cancelled mid-init."""
@@ -780,7 +994,6 @@ def _prepare(*, request: Any) -> Any:
assert response.status == ScenarioRunState.CANCELLED
execute.assert_not_called()
assert "cancelled-during-init" not in service._active_tasks
- assert service._run_semaphore._value == _DEFAULT_MAX_CONCURRENT_RUNS
async def test_start_run_resumes_a_run_that_was_already_cancelled(self, mock_all_registries) -> None:
"""Resuming a cancelled run is deliberate, so its starting state must not look like a cancel."""
@@ -828,7 +1041,6 @@ def _prepare(*, request: Any) -> Any:
assert response.status == ScenarioRunState.CANCELLED
execute.assert_not_called()
- assert service._run_semaphore._value == _DEFAULT_MAX_CONCURRENT_RUNS
async def test_start_run_does_not_read_a_header_for_a_fresh_run(self, mock_all_registries) -> None:
"""A fresh run has no stored state, so it must not pay for an extra query."""
@@ -857,7 +1069,6 @@ def _failing_prepare(*, request: Any) -> Any:
await service.start_run_async(request=_make_request())
assert "cancelled" not in caplog.text.lower()
- assert service._run_semaphore._value == _DEFAULT_MAX_CONCURRENT_RUNS
async def test_start_run_cleanup_failure_still_propagates_cancellation(self, mock_all_registries) -> None:
"""Cleanup runs inline on this path, so it must not replace the CancelledError."""
@@ -1042,8 +1253,8 @@ async def _link(depth: int) -> None:
assert time.monotonic() - started < 3
- async def test_start_run_releases_semaphore_when_initialization_leaks_a_task(self, mock_all_registries) -> None:
- """Failing the preparation must not strand the permit it was holding."""
+ async def test_start_run_fails_when_initialization_leaks_a_task(self, mock_all_registries) -> None:
+ """A preparation with leaked tasks must fail before scheduling."""
service = ScenarioRunService()
async def _leaky_prepare(*, request: Any) -> Any:
@@ -1057,8 +1268,6 @@ async def _leaky_prepare(*, request: Any) -> Any:
with pytest.raises(RuntimeError, match="left background tasks"):
await service.start_run_async(request=_make_request())
- assert service._run_semaphore._value == _DEFAULT_MAX_CONCURRENT_RUNS
-
def test_prepare_run_blocking_is_quiet_when_initialization_is_self_contained(
self, mock_all_registries, caplog
) -> None:
@@ -1074,8 +1283,8 @@ async def _clean_prepare(*, request: Any) -> Any:
assert "left background tasks" not in caplog.text
- async def test_start_run_holds_semaphore_until_abandoned_prepare_finishes(self, mock_all_registries) -> None:
- """A cancelled start must not free capacity while its worker thread is still initializing."""
+ async def test_start_run_serializes_after_abandoned_prepare(self, mock_all_registries) -> None:
+ """A cancelled start must leave its worker isolated until initialization finishes."""
service = ScenarioRunService()
finished = threading.Event()
@@ -1091,17 +1300,13 @@ def _slow_prepare(*, request: Any) -> Any:
with pytest.raises(asyncio.CancelledError):
await task
- # The worker thread cannot be killed, so admitting another run here would let
- # two initializations share one permit.
assert not finished.is_set()
- assert service._run_semaphore._value == _DEFAULT_MAX_CONCURRENT_RUNS - 1
await asyncio.sleep(1.0)
assert finished.is_set()
- assert service._run_semaphore._value == _DEFAULT_MAX_CONCURRENT_RUNS
- async def test_start_run_releases_semaphore_when_prepare_fails(self, mock_all_registries) -> None:
- """CancelledError is a BaseException, so it needs an explicit release path."""
+ async def test_start_run_propagates_prepare_failure(self, mock_all_registries) -> None:
+ """Preparation failures must reach the caller."""
service = ScenarioRunService()
def _failing_prepare(*, request: Any) -> Any:
@@ -1111,10 +1316,8 @@ def _failing_prepare(*, request: Any) -> Any:
with pytest.raises(ValueError, match="boom"):
await service.start_run_async(request=_make_request())
- assert service._run_semaphore._value == _DEFAULT_MAX_CONCURRENT_RUNS
-
- async def test_start_run_releases_semaphore_when_result_id_missing(self, mock_all_registries) -> None:
- """The missing scenario_result_id check used to sit outside the try block."""
+ async def test_start_run_rejects_missing_result_id(self, mock_all_registries) -> None:
+ """A prepared scenario must provide a result ID."""
service = ScenarioRunService()
scenario_instance = mock_all_registries["scenario_instance"]
scenario_instance._scenario_result_id = None
@@ -1122,21 +1325,25 @@ async def test_start_run_releases_semaphore_when_result_id_missing(self, mock_al
with pytest.raises(ValueError, match="did not produce a scenario_result_id"):
await service.start_run_async(request=_make_request())
- assert service._run_semaphore._value == _DEFAULT_MAX_CONCURRENT_RUNS
+ async def test_start_run_rejects_unpersisted_initialized_scenario(self, mock_all_registries) -> None:
+ service = ScenarioRunService()
+ mock_all_registries["memory"].get_scenario_results.return_value = []
+
+ with pytest.raises(RuntimeError, match="was not persisted during initialization"):
+ await service.start_run_async(request=_make_request())
async def test_start_run_cleans_up_when_response_lookup_fails(self, mock_all_registries) -> None:
- """A response failure must not strand a permit or leave an active-task entry."""
+ """A response failure must not leave an active-task entry."""
service = ScenarioRunService()
- with patch.object(service, "get_run", return_value=None):
+ with patch.object(service, "_build_response", return_value=None):
with pytest.raises(RuntimeError, match="not found in the database"):
await service.start_run_async(request=_make_request())
- assert service._run_semaphore._value == _DEFAULT_MAX_CONCURRENT_RUNS
assert service._active_tasks == {}
- async def test_start_run_releases_semaphore_exactly_once_on_success(self, mock_all_registries) -> None:
- """The background task owns the permit after handoff, so it is not double-released."""
+ async def test_start_run_executes_successfully(self, mock_all_registries) -> None:
+ """The scheduler starts an initialized scenario."""
service = ScenarioRunService()
released = asyncio.Event()
@@ -1149,8 +1356,6 @@ async def _run_async() -> None:
await asyncio.wait_for(released.wait(), timeout=5)
await asyncio.sleep(0)
- assert service._run_semaphore._value == _DEFAULT_MAX_CONCURRENT_RUNS
-
class TestScenarioRunServiceGetRun:
"""Tests for ScenarioRunService.get_run."""
@@ -1608,6 +1813,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
@@ -1630,6 +1836,290 @@ 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_rejects_empty_nonterminal_page_with_more_rows(self, mock_memory) -> None:
+ mock_memory.get_scenario_run_state_page.return_value = ([], True)
+
+ with pytest.raises(RuntimeError, match="another page without returning a cursor row"):
+ await ScenarioRunService().reconcile_interrupted_runs_async()
+
+ 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.try_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
+ + mock_memory.try_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)
+
+ async def test_shutdown_waits_for_preparation_before_terminalizing_run(self, mock_all_registries) -> None:
+ mock_memory = mock_all_registries["memory"]
+ records = {
+ "active": _make_db_scenario_result(
+ result_id="active",
+ run_state=ScenarioRunState.IN_PROGRESS,
+ ),
+ "preparing": _make_db_scenario_result(
+ result_id="preparing",
+ run_state=ScenarioRunState.CREATED,
+ ),
+ }
+ active_started = asyncio.Event()
+ preparation_started = threading.Event()
+ release_preparation = threading.Event()
+ shutdown_started = asyncio.Event()
+
+ async def _run_active() -> None:
+ active_started.set()
+ await asyncio.Event().wait()
+
+ active_scenario = MagicMock()
+ active_scenario.run_async = AsyncMock(side_effect=_run_active)
+ prepared_scenario = MagicMock()
+ prepared_scenario._scenario_result_id = "preparing"
+ prepared_scenario.run_async = AsyncMock()
+
+ 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
+
+ def _try_update_state(
+ *,
+ scenario_result_id: str,
+ expected_states: set[ScenarioRunState],
+ scenario_run_state: ScenarioRunState,
+ **_: object,
+ ) -> bool:
+ record = records[scenario_result_id]
+ if record.scenario_run_state not in expected_states:
+ return False
+ record.scenario_run_state = scenario_run_state
+ return True
+
+ def _prepare(*, request: Any) -> MagicMock:
+ preparation_started.set()
+ if not release_preparation.wait(timeout=5):
+ raise TimeoutError("Test did not release scenario preparation.")
+ return prepared_scenario
+
+ async def _shutdown() -> None:
+ shutdown_started.set()
+ await service.shutdown_async()
+
+ mock_memory.get_scenario_results.side_effect = _get_results
+ mock_memory.update_scenario_run_state.side_effect = _update_state
+ mock_memory.try_update_scenario_run_state.side_effect = _try_update_state
+ mock_memory.update_scenario_run_state_and_metadata_fields.side_effect = _update_state
+
+ service = ScenarioRunService()
+ active = _svc_mod._ActiveTask(
+ scenario_result_id="active",
+ scenario=active_scenario,
+ )
+ service._active_scenario_result_id = "active"
+ service._active_tasks["active"] = active
+ active.task = asyncio.create_task(service._execute_run_async(scenario_result_id="active"))
+ await active_started.wait()
+
+ with patch.object(service, "_prepare_run_blocking", _prepare):
+ start_task = asyncio.create_task(service.start_run_async(request=_make_request()))
+ assert await asyncio.to_thread(preparation_started.wait, 5)
+ shutdown_task = asyncio.create_task(_shutdown())
+ await shutdown_started.wait()
+
+ assert not shutdown_task.done()
+ assert records["preparing"].scenario_run_state == ScenarioRunState.CREATED
+ release_preparation.set()
+ await asyncio.wait_for(asyncio.gather(start_task, shutdown_task), timeout=5)
+
+ assert records["active"].scenario_run_state == ScenarioRunState.FAILED
+ assert records["preparing"].scenario_run_state == ScenarioRunState.FAILED
+ assert service.get_queue_snapshot().active is None
+ assert service.get_queue_snapshot().queued == []
+ prepared_scenario.run_async.assert_not_awaited()
+
+ async def test_shutdown_reports_all_persistence_failures_and_clears_scheduler(self, mock_all_registries) -> None:
+ service = ScenarioRunService()
+ completed_task = MagicMock(spec=asyncio.Task)
+ completed_task.done.return_value = True
+ service._active_scenario_result_id = "active"
+ service._active_tasks["active"] = _svc_mod._ActiveTask(
+ scenario_result_id="active",
+ task=completed_task,
+ scenario=MagicMock(),
+ )
+ service._queued_runs.append(
+ _svc_mod._ActiveTask(
+ scenario_result_id="queued",
+ scenario=MagicMock(),
+ )
+ )
+ mock_all_registries["memory"].update_scenario_run_state.side_effect = [
+ RuntimeError("queued persistence failed"),
+ RuntimeError("active persistence failed"),
+ ]
+
+ with pytest.raises(ExceptionGroup) as exc_info:
+ await service.shutdown_async()
+
+ assert [str(error) for error in exc_info.value.exceptions] == [
+ "queued persistence failed",
+ "active persistence failed",
+ ]
+ assert service.get_queue_snapshot().active is None
+ assert service.get_queue_snapshot().queued == []
+ assert service._active_tasks == {}
+
+ def test_overload_summaries_group_429_and_5xx_by_role_without_false_positives(self, mock_memory) -> None:
+ now = datetime(2025, 1, 1, tzinfo=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),
+ MagicMock(component_role="objective_target", status_code=429, timestamp="invalid"),
+ ]
+
+ 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"]
@@ -1756,46 +2246,127 @@ async def test_execute_run_completes_successfully(self, mock_all_registries) ->
mock_scenario_result.creation_time = datetime(2025, 1, 1, tzinfo=UTC)
mock_scenario_result.completion_time = datetime(2025, 1, 1, 0, 5, tzinfo=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
+
+ async def test_execute_run_retains_richer_terminal_error_from_scenario(self, mock_all_registries) -> None:
+ service = ScenarioRunService()
+ mock_instance = mock_all_registries["scenario_instance"]
+ mock_memory = mock_all_registries["memory"]
+ persisted = mock_all_registries["db_result"]
+ execution_started = asyncio.Event()
+ release_execution = asyncio.Event()
+
+ def _update_state_and_metadata(
+ *,
+ scenario_run_state: ScenarioRunState,
+ **_: object,
+ ) -> None:
+ persisted.scenario_run_state = scenario_run_state
+
+ def _try_update_state(
+ *,
+ expected_states: set[ScenarioRunState],
+ scenario_run_state: ScenarioRunState,
+ error_message: str | None = None,
+ error_type: str | None = None,
+ **_: object,
+ ) -> bool:
+ if persisted.scenario_run_state not in expected_states:
+ return False
+ persisted.scenario_run_state = scenario_run_state
+ persisted.error_message = error_message
+ persisted.error_type = error_type
+ return True
+
+ async def _run() -> None:
+ execution_started.set()
+ await release_execution.wait()
+ persisted.scenario_run_state = ScenarioRunState.FAILED
+ persisted.error_message = "The target rejected the request body."
+ persisted.error_type = "BadRequestError"
+ raise RuntimeError("One or more attacks failed.")
+
+ mock_memory.update_scenario_run_state_and_metadata_fields.side_effect = _update_state_and_metadata
+ mock_memory.try_update_scenario_run_state.side_effect = _try_update_state
+ mock_instance.run_async = AsyncMock(side_effect=_run)
+
+ response = await service.start_run_async(request=_make_request())
+ await execution_started.wait()
+ 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
+
+ fetched = service.get_run(scenario_result_id=response.scenario_result_id)
+ assert fetched is not None
+ assert fetched.status == ScenarioRunState.FAILED
+ assert fetched.error == "The target rejected the request body."
+ assert fetched.error_type == "BadRequestError"
+ mock_memory.update_scenario_run_state.assert_not_called()
+ mock_memory.try_update_scenario_run_state.assert_called_once_with(
+ scenario_result_id=response.scenario_result_id,
+ expected_states={ScenarioRunState.CREATED, ScenarioRunState.IN_PROGRESS},
+ scenario_run_state=ScenarioRunState.FAILED,
+ error_message="One or more attacks failed.",
+ error_type="RuntimeError",
+ )
class TestScenarioRunServiceGetResults:
@@ -1960,6 +2531,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
@@ -2234,6 +2828,73 @@ def test_planned_progress_includes_latest_errors_in_success_rate_denominator(moc
assert summary.objective_achieved_rate == 50
+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=to_sha256(objective),
+ objective=objective,
+ )
+ ],
+ )
+ timestamp = datetime(2026, 8, 8, tzinfo=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": to_sha256(objective)}],
+ )
+ aggregate = ScenarioHistoryAggregate(
+ scenario_result_id=str(scenario_result.id),
+ unit_count=1,
+ completed_units=1,
+ successful_units=1,
+ error_attempts=2,
+ total_retries=5,
+ latest_attempt_timestamp=timestamp + timedelta(seconds=2),
+ atomic_attack_names=("attack",),
+ )
+ service = ScenarioRunService()
+
+ detail = service._build_response_from_db(scenario_result=scenario_result)
+ history = service._build_history_summary(
+ record=record,
+ atomic_groups=plan.atomic_groups,
+ aggregate=aggregate,
+ )
+ 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)])
@@ -2456,6 +3117,43 @@ def test_get_progress_cache_refreshes_identifier_enriched_after_insert(mock_memo
assert mock_memory.get_scenario_attack_result_deltas.call_args_list[1].kwargs["cursor"] is None
+def test_get_progress_exposes_persisted_started_at(mock_memory) -> None:
+ started_at = datetime(2026, 8, 8, 12, 30, tzinfo=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
+
+
+@pytest.mark.parametrize("started_at", ["not-a-timestamp", "2026-08-08T12:30:00"])
+def test_load_started_at_rejects_invalid_or_naive_timestamp(started_at: str) -> None:
+ scenario_result = make_scenario_result(
+ attack_results={},
+ metadata={_svc_mod.SCENARIO_RUN_STARTED_AT_METADATA_KEY: started_at},
+ )
+
+ assert ScenarioRunService._load_started_at(scenario_result=scenario_result) is None
+
+
def test_get_progress_refreshes_enriched_rows_across_storage_pages(mock_memory: MagicMock) -> None:
attack_identifier = ComponentIdentifier(class_name="PromptSendingAttack", class_module="tests")
unenriched_identifier = AtomicAttackIdentifier.build(attack_identifier=attack_identifier)
diff --git a/tests/unit/backend/test_scenario_service.py b/tests/unit/backend/test_scenario_service.py
index 226b18a7de..f3f032901f 100644
--- a/tests/unit/backend/test_scenario_service.py
+++ b/tests/unit/backend/test_scenario_service.py
@@ -833,12 +833,15 @@ 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 = ScenarioRunSizeEstimate(
estimated_attack_count=12,
components=[ScenarioRunSizeComponent(label="Configured Jailbreak", count=12)],
)
- with patch("pyrit.backend.routes.scenarios.get_scenario_service") as mock_get_service:
+ 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
@@ -847,7 +850,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,
@@ -859,11 +862,12 @@ def test_estimate_scenario_returns_configured_projection(self, client: TestClien
assert response.json()["estimated_attack_count"] == 12
request = mock_service.estimate_scenario_run_size_async.await_args.kwargs["request"]
assert request.techniques == ["prompt_sending"]
- assert request.include_baseline is True
+ 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/infra/test_bicep_topology.py b/tests/unit/infra/test_bicep_topology.py
index e1dc524421..624ee3331b 100644
--- a/tests/unit/infra/test_bicep_topology.py
+++ b/tests/unit/infra/test_bicep_topology.py
@@ -154,6 +154,7 @@ def test_application_reads_existing_state_and_preserves_authentication(self) ->
for name in ("containerImage", "existingManagedIdentityResourceId"):
assert template["parameters"][name]["minLength"] == 1
assert "defaultValue" not in template["parameters"][name]
+ assert template["parameters"]["maxReplicas"]["allowedValues"] == [1]
assert "fail(" in template["variables"]["validatedAllowedGroupObjectIds"]
assert "fail(" in template["variables"]["validatedAdminGroupObjectId"]
assert "trim(" in template["variables"]["normalizedAllowedGroupObjectIds"]
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 2f58edfef9..b201c24097 100644
--- a/tests/unit/memory/memory_interface/test_interface_scenario_history.py
+++ b/tests/unit/memory/memory_interface/test_interface_scenario_history.py
@@ -5,15 +5,18 @@
import json
import uuid
+from contextlib import closing
from datetime import UTC, datetime, timedelta
-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.common.utils import to_sha256
-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,
@@ -25,6 +28,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"),
[
@@ -182,6 +202,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)),
@@ -198,6 +223,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)),
@@ -233,6 +259,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 = (
@@ -525,6 +552,143 @@ def test_history_aggregates_fill_zero_for_runs_without_attempts(sqlite_instance:
assert aggregates[scenario_result_id].latest_attempt_timestamp is None
+def test_legacy_label_hook_is_constructible_and_composes_multi_value_semantics(
+ sqlite_instance: MemoryInterface,
+) -> None:
+ timestamp = datetime(2026, 8, 7, tzinfo=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=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
+
+
+@pytest.mark.parametrize("limit", [0, 501])
+def test_nonterminal_state_projection_rejects_out_of_range_limit(
+ sqlite_instance: MemoryInterface,
+ limit: int,
+) -> None:
+ with pytest.raises(ValueError, match="between 1 and 500"):
+ sqlite_instance.get_scenario_run_state_page(
+ states=[ScenarioRunState.QUEUED],
+ limit=limit,
+ )
+
+
+def test_state_and_metadata_update_persists_scheduler_start_atomically(
+ sqlite_instance: MemoryInterface,
+) -> None:
+ timestamp = datetime(2026, 8, 7, tzinfo=UTC)
+ scenario = _make_scenario(
+ result_id=uuid.UUID(int=43),
+ timestamp=timestamp,
+ name="Scheduled",
+ state=ScenarioRunState.CREATED,
+ labels={},
+ registry_name="registered.scenario",
+ )
+ sqlite_instance.add_scenario_results_to_memory(scenario_results=[scenario])
+
+ sqlite_instance.update_scenario_run_state_and_metadata_fields(
+ scenario_result_id=str(scenario.id),
+ scenario_run_state=ScenarioRunState.IN_PROGRESS,
+ metadata_fields={"started_at": timestamp.isoformat()},
+ )
+
+ stored = sqlite_instance.get_scenario_result_header(scenario_result_id=str(scenario.id))
+ assert stored is not None
+ assert stored.scenario_run_state is ScenarioRunState.IN_PROGRESS
+ assert stored.metadata["started_at"] == timestamp.isoformat()
+ assert SCENARIO_RUN_PLAN_METADATA_KEY in stored.metadata
+
+
+def test_metadata_field_update_rejects_unknown_run(sqlite_instance: MemoryInterface) -> None:
+ with pytest.raises(ValueError, match="not found in memory"):
+ sqlite_instance.update_scenario_metadata_fields(
+ scenario_result_id=str(uuid.UUID(int=44)),
+ fields={"started_at": datetime(2026, 8, 7, tzinfo=UTC).isoformat()},
+ )
+
+
+def test_started_at_parser_rejects_malformed_timestamp() -> None:
+ assert MemoryInterface._parse_scenario_started_at(raw_value="not-a-timestamp") is None
+
+
+def test_default_started_at_projection_is_null() -> None:
+ expression = MemoryInterface._get_scenario_started_at_expression(MagicMock())
+
+ assert expression.compile().params == {"param_1": None}
+
+
def test_unique_scenario_labels_are_grouped_for_filter_options(sqlite_instance: MemoryInterface) -> None:
timestamp = datetime(2026, 8, 7, tzinfo=UTC)
scenarios = [
diff --git a/tests/unit/memory/test_azure_sql_memory.py b/tests/unit/memory/test_azure_sql_memory.py
index 41d6ce5bbe..b0794edeed 100644
--- a/tests/unit/memory/test_azure_sql_memory.py
+++ b/tests/unit/memory/test_azure_sql_memory.py
@@ -460,7 +460,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(
@@ -489,6 +489,32 @@ def test_scenario_history_conditions_bind_or_within_label_and_registry_values(
assert "scenario_registry_name_1" in combined_statement.compile().params
+def test_scenario_history_legacy_label_condition_binds_each_value(
+ memory_interface: AzureSQLMemory,
+) -> None:
+ condition = memory_interface._get_scenario_result_label_condition(
+ labels={"team.name": "alice", "operation": "nightly"}
+ )
+
+ assert condition.compile().params == {
+ "scenario_label_path_0": '$."team.name"',
+ "scenario_label_value_0": "alice",
+ "scenario_label_path_1": '$."operation"',
+ "scenario_label_value_1": "nightly",
+ }
+ assert " AND " in str(condition)
+
+
+def test_scenario_history_started_at_uses_sql_server_json_value(
+ memory_interface: AzureSQLMemory,
+) -> None:
+ expression = memory_interface._get_scenario_started_at_expression()
+
+ compiled = select(expression).compile()
+ assert "json_value" in str(compiled).lower()
+ assert "$.started_at" in compiled.params.values()
+
+
def test_scenario_history_seed_projection_defaults_to_empty_json(memory_interface: AzureSQLMemory) -> None:
"""The SQL Server seed projection returns an empty JSON array for runs without seed groups."""
_, _, seed_projection = memory_interface._get_scenario_history_plan_expressions()
diff --git a/tests/unit/registry/test_scenario_registry.py b/tests/unit/registry/test_scenario_registry.py
index a78fd5ab6d..649ef9beb7 100644
--- a/tests/unit/registry/test_scenario_registry.py
+++ b/tests/unit/registry/test_scenario_registry.py
@@ -132,6 +132,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,
)
@@ -139,6 +140,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 e47d4e7451..3f3b08fc89 100644
--- a/tests/unit/scenario/core/test_scenario.py
+++ b/tests/unit/scenario/core/test_scenario.py
@@ -270,6 +270,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)
@@ -277,6 +278,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 = [
@@ -357,6 +359,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(