=> {
+ if (submitting.current) return
+ submitting.current = true
+ setDeleting(true)
+ setError('')
+ try {
+ await operationsApi.delete(operation.id)
+ setOpen(false)
+ onDeleted()
+ } catch (cause: unknown) {
+ const apiError = toApiError(cause)
+ if (apiError.status === 409) onInUse()
+ setError(apiError.status === 409 ? 'This operation is no longer empty, so it can’t be deleted.' : apiError.detail)
+ } finally {
+ submitting.current = false
+ setDeleting(false)
+ }
+ }
+
+ return (
+
+ )
+}
diff --git a/frontend/src/components/Operations/OperationDetailPage.test.tsx b/frontend/src/components/Operations/OperationDetailPage.test.tsx
index 68c2f0deb3..5d2ec2cb30 100644
--- a/frontend/src/components/Operations/OperationDetailPage.test.tsx
+++ b/frontend/src/components/Operations/OperationDetailPage.test.tsx
@@ -4,12 +4,12 @@ import { FluentProvider, webLightTheme } from '@fluentui/react-components'
import { MemoryRouter, Route, Routes } from 'react-router'
import { attacksApi, operationsApi, scenariosApi } from '@/services/api'
-import type { Finding, FindingListItem, Operation } from '@/types'
+import type { Finding, FindingListItem, Operation, OperationDetail } from '@/types'
import OperationDetailPage from './OperationDetailPage'
jest.mock('@/services/api', () => ({
operationsApi: {
- get: jest.fn(), update: jest.fn(), listFindings: jest.fn(), createFinding: jest.fn(),
+ get: jest.fn(), update: jest.fn(), delete: jest.fn(), listFindings: jest.fn(), createFinding: jest.fn(),
updateFinding: jest.fn(), deleteFinding: jest.fn(),
listFindingEvidence: jest.fn(), detachFindingEvidence: jest.fn(),
getFindingOptions: jest.fn(),
@@ -18,20 +18,24 @@ jest.mock('@/services/api', () => ({
scenariosApi: { listRuns: jest.fn() },
}))
-const OPERATION: Operation = { id: 'op-1', name: 'Red team / α%', created_at: '2026-10-06T20:00:00Z' }
+const SAVED: Operation = { id: 'op-1', name: 'Red team / α%', created_at: '2026-10-06T20:00:00Z' }
+const OPERATION: OperationDetail = { ...SAVED, usage: { findings: 0, attacks: 0, scenario_runs: 0 } }
const FINDING: FindingListItem = {
id: 'finding-1', operation_id: 'op-1', title: 'Human assessment',
description: '', severity: 'informational', created_at: '2026-10-06T20:00:00Z', evidence_count: 0,
}
const EMPTY = { items: [], has_more: false, next_offset: null }
const NOT_FOUND = { isAxiosError: true, response: { status: 404, data: { detail: 'missing' } } }
+const conflict = (detail: object) => ({ isAxiosError: true, response: { status: 409, data: { detail } } })
-function renderPage(route = '/operations/op-1'): void {
+function renderPage(route = '/operations/op-1', onOperationChanged?: () => void): void {
render(
- } />
+ Operations list} />
+ } />
,
@@ -52,25 +56,130 @@ beforeEach(() => {
jest.mocked(scenariosApi.listRuns).mockResolvedValue({ items: [], pagination: { limit: 5, has_more: false } })
})
-it('adds and edits the operation description, keeping the draft after a failed save', async () => {
+it('renames the operation and edits its description, keeping the draft after a failed or conflicting save', async () => {
const user = userEvent.setup()
- renderPage()
- await user.click(await screen.findByRole('button', { name: 'Add description' }))
- let dialog = screen.getByRole('dialog')
+ const onOperationChanged = jest.fn()
+ renderPage(undefined, onOperationChanged)
+ await user.click(await screen.findByRole('button', { name: 'Edit' }))
+ let dialog = screen.getByRole('dialog', { name: 'Edit operation' })
+ const name = within(dialog).getByRole('textbox', { name: 'Name' })
+ expect(name).toHaveValue(OPERATION.name)
expect(within(dialog).getByRole('textbox', { name: 'Description' })).toHaveValue('')
+ await user.clear(name)
+ await user.type(name, 'Engagement')
await user.type(within(dialog).getByRole('textbox', { name: 'Description' }), 'Web app only')
jest.mocked(operationsApi.update).mockRejectedValueOnce(new Error('Storage unavailable'))
- await user.click(within(dialog).getByRole('button', { name: 'Save description' }))
- expect(await within(dialog).findByText(/Could not save description: Storage unavailable/)).toBeInTheDocument()
- jest.mocked(operationsApi.update).mockResolvedValueOnce({ ...OPERATION, description: 'Web app only' })
- await user.click(within(dialog).getByRole('button', { name: 'Save description' }))
- expect(await screen.findByText('Web app only')).toBeInTheDocument()
- expect(screen.getByText('Operation description saved.')).toBeInTheDocument()
- expect(operationsApi.update).toHaveBeenLastCalledWith('op-1', { description: 'Web app only' })
+ await user.click(within(dialog).getByRole('button', { name: 'Save operation' }))
+ expect(await within(dialog).findByText(/Could not save operation: Storage unavailable/)).toBeInTheDocument()
+ jest.mocked(operationsApi.update).mockRejectedValueOnce(conflict({
+ message: 'taken', operation: { id: 'op-2', name: 'ENGAGEMENT', created_at: SAVED.created_at },
+ }))
+ await user.click(within(dialog).getByRole('button', { name: 'Save operation' }))
+ expect(await within(dialog).findByText(/Another operation is already named “ENGAGEMENT”/)).toBeInTheDocument()
+ expect(name).toHaveValue('Engagement')
+ await user.type(name, ' 2')
+ expect(within(dialog).queryByText(/already named/)).not.toBeInTheDocument()
+ jest.mocked(operationsApi.update).mockResolvedValueOnce({ ...SAVED, name: 'Engagement 2', description: 'Web app only' })
+ await user.click(within(dialog).getByRole('button', { name: 'Save operation' }))
+ expect(await screen.findByRole('heading', { level: 1, name: 'Engagement 2' })).toBeInTheDocument()
+ expect(screen.getByText('Web app only')).toBeInTheDocument()
+ expect(screen.getByText('Operation updated.')).toBeInTheDocument()
+ expect(operationsApi.update).toHaveBeenLastCalledWith('op-1', { name: 'Engagement 2', description: 'Web app only' })
+ expect(onOperationChanged).toHaveBeenCalledTimes(1)
+ expect(screen.getByRole('button', { name: 'Delete' })).not.toHaveAttribute('aria-disabled')
- await user.click(await screen.findByRole('button', { name: 'Edit description' }))
- dialog = screen.getByRole('dialog')
+ await user.click(screen.getByRole('button', { name: 'Edit' }))
+ dialog = screen.getByRole('dialog', { name: 'Edit operation' })
+ expect(within(dialog).getByRole('textbox', { name: 'Name' })).toHaveValue('Engagement 2')
expect(within(dialog).getByRole('textbox', { name: 'Description' })).toHaveValue('Web app only')
+ await user.clear(within(dialog).getByRole('textbox', { name: 'Name' }))
+ expect(within(dialog).getByRole('button', { name: 'Save operation' })).toBeDisabled()
+})
+
+it('sends only the fields that were edited', async () => {
+ const user = userEvent.setup()
+ jest.mocked(operationsApi.update).mockResolvedValue({ ...SAVED, description: 'Web app only' })
+ renderPage()
+ await user.click(await screen.findByRole('button', { name: 'Edit' }))
+ const dialog = screen.getByRole('dialog', { name: 'Edit operation' })
+ await user.type(within(dialog).getByRole('textbox', { name: 'Description' }), 'Web app only')
+ await user.click(within(dialog).getByRole('button', { name: 'Save operation' }))
+ expect(await screen.findByText('Web app only')).toBeInTheDocument()
+ expect(operationsApi.update).toHaveBeenCalledWith('op-1', { description: 'Web app only' })
+})
+
+it('deletes an empty operation after confirmation and returns to the operations list', async () => {
+ const user = userEvent.setup()
+ const onOperationChanged = jest.fn()
+ jest.mocked(operationsApi.delete).mockResolvedValue(undefined)
+ renderPage(undefined, onOperationChanged)
+ const remove = await screen.findByRole('button', { name: 'Delete' })
+ expect(remove).toHaveAccessibleDescription('Delete this empty operation.')
+ await user.click(remove)
+ let dialog = screen.getByRole('dialog', { name: 'Delete operation?' })
+ const confirmation = within(dialog).getByText(/cannot be undone/i)
+ expect(confirmation).toHaveTextContent(OPERATION.name)
+ expect(confirmation).toHaveTextContent(/No saved findings, attacks, or scanner runs use it/)
+ expect(confirmation).toHaveTextContent(/SDK attacks still running under it will save their results without an operation/)
+ await user.click(within(dialog).getByRole('button', { name: 'Cancel' }))
+ await waitFor(() => expect(screen.queryByRole('dialog')).not.toBeInTheDocument())
+ expect(operationsApi.delete).not.toHaveBeenCalled()
+ await user.click(remove)
+ dialog = screen.getByRole('dialog', { name: 'Delete operation?' })
+ await user.click(within(dialog).getByRole('button', { name: 'Delete operation' }))
+ // JSDOM can leave Tabster's modal aria-hidden after navigation; the seeded browser test checks the accessible page.
+ expect(await screen.findByText('Operations list')).toBeInTheDocument()
+ expect(operationsApi.delete).toHaveBeenCalledWith('op-1')
+ expect(onOperationChanged).toHaveBeenCalledTimes(1)
+})
+
+it.each([
+ [{ findings: 3, attacks: 10, scenario_runs: 2 }, 'Has 3 findings and 12 linked attacks or runs.'],
+ [{ findings: 1, attacks: 0, scenario_runs: 0 }, 'Has 1 finding.'],
+ [{ findings: 0, attacks: 0, scenario_runs: 1 }, 'Has 1 linked attack or run.'],
+])('keeps an operation in use undeletable and explains why: %o', async (usage, reason) => {
+ const user = userEvent.setup()
+ jest.mocked(operationsApi.get).mockResolvedValue({ ...OPERATION, usage })
+ renderPage()
+ const remove = await screen.findByRole('button', { name: 'Delete' })
+ expect(remove).toHaveAttribute('aria-disabled', 'true')
+ expect(remove).toHaveAccessibleDescription(`${reason} Only empty operations can be deleted.`)
+ remove.focus()
+ expect(remove).toHaveFocus()
+ await user.keyboard('{Enter}')
+ await user.click(remove)
+ expect(screen.queryByRole('dialog')).not.toBeInTheDocument()
+ expect(operationsApi.delete).not.toHaveBeenCalled()
+})
+
+it('refetches the counts and disables deletion when the operation came into use after loading', async () => {
+ const user = userEvent.setup()
+ jest.mocked(operationsApi.delete).mockRejectedValueOnce(conflict({ message: 'in use' }))
+ renderPage()
+ await user.click(await screen.findByRole('button', { name: 'Delete' }))
+ const dialog = screen.getByRole('dialog', { name: 'Delete operation?' })
+ jest.mocked(operationsApi.get).mockResolvedValue({ ...OPERATION, usage: { findings: 0, attacks: 2, scenario_runs: 0 } })
+ await user.click(within(dialog).getByRole('button', { name: 'Delete operation' }))
+ expect(await within(dialog).findByText(/no longer empty/)).toBeInTheDocument()
+ await waitFor(() => {
+ expect(within(dialog).getByRole('button', { name: 'Delete operation' })).toHaveAttribute('aria-disabled', 'true')
+ })
+ await user.click(within(dialog).getByRole('button', { name: 'Cancel' }))
+ const remove = await screen.findByRole('button', { name: 'Delete' })
+ expect(remove).toHaveAttribute('aria-disabled', 'true')
+ expect(remove).toHaveAccessibleDescription('Has 2 linked attacks or runs. Only empty operations can be deleted.')
+ expect(screen.getByRole('heading', { level: 1, name: OPERATION.name })).toBeInTheDocument()
+})
+
+it('reports other deletion failures in the confirmation dialog', async () => {
+ const user = userEvent.setup()
+ jest.mocked(operationsApi.delete).mockRejectedValueOnce(new Error('Storage unavailable'))
+ renderPage()
+ await user.click(await screen.findByRole('button', { name: 'Delete' }))
+ const dialog = screen.getByRole('dialog', { name: 'Delete operation?' })
+ await user.click(within(dialog).getByRole('button', { name: 'Delete operation' }))
+ expect(await within(dialog).findByText('Storage unavailable')).toBeInTheDocument()
+ expect(within(dialog).getByRole('button', { name: 'Delete operation' })).not.toHaveAttribute('aria-disabled')
})
it('refreshes evidence counts without collapsing the expanded list after detachment', async () => {
@@ -323,6 +432,30 @@ it('requires delete confirmation and refreshes the first page after deleting its
await waitFor(() => expect(screen.getByRole('button', { name: 'New finding' })).toHaveFocus())
})
+it('refreshes operation usage after findings change so an emptied operation can be deleted', async () => {
+ const user = userEvent.setup()
+ const usage = { findings: 1, attacks: 0, scenario_runs: 0 }
+ jest.mocked(operationsApi.get).mockResolvedValue({ ...OPERATION, usage })
+ jest.mocked(operationsApi.listFindings).mockResolvedValue({ ...EMPTY, items: [FINDING] })
+ renderPage()
+ expect(await screen.findByRole('button', { name: 'Delete' })).toHaveAttribute('aria-disabled', 'true')
+ await user.click(screen.getByRole('button', { name: `Delete finding: ${FINDING.title}` }))
+ jest.mocked(operationsApi.get).mockResolvedValue(OPERATION)
+ jest.mocked(operationsApi.listFindings).mockResolvedValue(EMPTY)
+ await user.click(within(screen.getByRole('dialog')).getByRole('button', { name: 'Delete finding' }))
+ await waitFor(() => expect(screen.getByRole('button', { name: 'Delete' })).not.toHaveAttribute('aria-disabled'))
+ expect(operationsApi.get).toHaveBeenCalledTimes(2)
+
+ jest.mocked(operationsApi.get).mockRejectedValue(new Error('Busy'))
+ await user.click(screen.getByRole('button', { name: 'New finding' }))
+ const dialog = screen.getByRole('dialog')
+ await user.type(within(dialog).getByRole('textbox', { name: 'Title' }), FINDING.title)
+ await user.click(within(dialog).getByRole('button', { name: 'Save finding' }))
+ await waitFor(() => expect(operationsApi.get).toHaveBeenCalledTimes(3))
+ expect(await screen.findByRole('heading', { level: 1, name: OPERATION.name })).toBeInTheDocument()
+ expect(screen.queryByText(/Could not load operation/)).not.toBeInTheDocument()
+})
+
it('keeps deletion errors visible and prevents duplicate confirmations', async () => {
const user = userEvent.setup()
jest.mocked(operationsApi.listFindings).mockResolvedValue({ ...EMPTY, items: [FINDING] })
diff --git a/frontend/src/components/Operations/OperationDetailPage.tsx b/frontend/src/components/Operations/OperationDetailPage.tsx
index 9cfaf01164..9f777d6b66 100644
--- a/frontend/src/components/Operations/OperationDetailPage.tsx
+++ b/frontend/src/components/Operations/OperationDetailPage.tsx
@@ -1,5 +1,5 @@
import { useEffect, useRef, useState } from 'react'
-import { Link, useParams, useSearchParams } from 'react-router'
+import { Link, useNavigate, useParams, useSearchParams } from 'react-router'
import {
Button, Dialog, DialogBody, DialogContent, DialogSurface, DialogTitle,
MessageBar, MessageBarBody, Spinner, Text,
@@ -9,27 +9,34 @@ import { ArrowLeftRegular } from '@fluentui/react-icons'
import { DEFAULT_HISTORY_FILTERS, filtersToSearchParams } from '@/components/History/historyFilters'
import { operationsApi } from '@/services/api'
import { toApiError } from '@/services/errors'
-import type { Finding, FindingCreate, FindingListItem, FindingListResponse, Operation } from '@/types'
+import type { Finding, FindingCreate, FindingListItem, FindingListResponse, OperationDetail } from '@/types'
import FindingDialog, { type FindingSaveAttempt } from './FindingDialog'
import FindingEvidenceList from './FindingEvidenceList'
import FindingSeverityBadge from './FindingSeverityBadge'
-import OperationDescriptionDialog from './OperationDescriptionDialog'
+import OperationDeleteDialog from './OperationDeleteDialog'
+import OperationEditDialog from './OperationEditDialog'
import { useOperationsStyles } from './Operations.styles'
const PAGE_SIZE = 20
type OperationState =
- | { key: string; status: 'loaded'; operation: Operation }
+ | { key: string; status: 'loaded'; operation: OperationDetail }
| { key: string; status: 'missing' }
| { key: string; status: 'error'; error: string }
-export default function OperationDetailPage() {
+interface OperationDetailPageProps {
+ /** Called after the Operation is renamed or deleted, so views holding it can refresh. */
+ onOperationChanged?: () => void
+}
+
+export default function OperationDetailPage({ onOperationChanged }: OperationDetailPageProps) {
const { operationId = '' } = useParams()
- return
+ return
}
-function OperationDetailContent({ operationId }: { operationId: string }) {
+function OperationDetailContent({ operationId, onOperationChanged }: { operationId: string } & OperationDetailPageProps) {
const styles = useOperationsStyles()
+ const navigate = useNavigate()
const [params, setParams] = useSearchParams()
const rawOffset = Number(params.get('offset') ?? 0)
const offset = Number.isSafeInteger(rawOffset) && rawOffset >= 0 ? rawOffset : 0
@@ -70,12 +77,13 @@ function OperationDetailContent({ operationId }: { operationId: string }) {
.catch((cause: unknown) => {
if (ignore) return
const apiError = toApiError(cause)
- setOperationState(apiError.status === 404
- ? { key: operationKey, status: 'missing' }
- : { key: operationKey, status: 'error', error: apiError.detail })
+ // A failed usage refresh after a finding change keeps the page; the server still checks before deleting.
+ setOperationState(state => state?.key === operationKey && state.status === 'loaded' ? state
+ : apiError.status === 404 ? { key: operationKey, status: 'missing' }
+ : { key: operationKey, status: 'error', error: apiError.detail })
})
return () => { ignore = true }
- }, [operationId, operationKey])
+ }, [operationId, operationKey, revision])
useEffect(() => {
let ignore = false
@@ -183,10 +191,17 @@ function OperationDetailContent({ operationId }: { operationId: string }) {
{current.operation.name}
- {
- setOperationState({ key: operationKey, status: 'loaded', operation })
- setNotice('Operation description saved.')
+ {
+ setOperationState({ key: operationKey, status: 'loaded', operation: { ...current.operation, ...operation } })
+ setNotice('Operation updated.')
+ onOperationChanged?.()
}} />
+ { setRevision(value => value + 1) }}
+ onDeleted={() => {
+ onOperationChanged?.()
+ void navigate('/operations')
+ }} />
diff --git a/frontend/src/services/api.ts b/frontend/src/services/api.ts
index 2bd21495cd..0122b335e1 100644
--- a/frontend/src/services/api.ts
+++ b/frontend/src/services/api.ts
@@ -76,6 +76,7 @@ import type {
FindingEvidenceListResponse,
Operation,
OperationCreate,
+ OperationDetail,
OperationListResponse,
OperationUpdate,
AttackAnalyticsQuery,
@@ -97,7 +98,7 @@ export const operationsApi = {
list: async (): Promise => {
return (await apiClient.get('/operations')).data
},
- get: async (operationId: string): Promise => {
+ get: async (operationId: string): Promise => {
return (await apiClient.get(operationPath(operationId))).data
},
create: async (request: OperationCreate): Promise => {
@@ -106,6 +107,9 @@ export const operationsApi = {
update: async (operationId: string, request: OperationUpdate): Promise => {
return (await apiClient.patch(operationPath(operationId), request)).data
},
+ delete: async (operationId: string): Promise => {
+ await apiClient.delete(operationPath(operationId))
+ },
listFindings: async (
operationId: string, params: { limit: number; offset: number; title?: string },
): Promise => {
diff --git a/frontend/src/types/index.ts b/frontend/src/types/index.ts
index 2946cb8bf5..1992df9af0 100644
--- a/frontend/src/types/index.ts
+++ b/frontend/src/types/index.ts
@@ -1799,9 +1799,10 @@ export interface OperationCreate {
description?: string
}
-/** The editable fields of a saved Operation; renaming is not supported yet. */
+/** The editable fields of a saved Operation. Omitted fields are unchanged. */
export interface OperationUpdate {
- description: string
+ name?: string
+ description?: string
}
export interface Operation extends OperationCreate {
@@ -1809,6 +1810,17 @@ export interface Operation extends OperationCreate {
created_at: string
}
+/** What refers to an Operation by ID. Only an Operation with no usage can be deleted. */
+export interface OperationUsage {
+ findings: number
+ attacks: number
+ scenario_runs: number
+}
+
+export interface OperationDetail extends Operation {
+ usage: OperationUsage
+}
+
export interface OperationListItem extends Operation {
finding_counts: Partial>
}
diff --git a/pyrit/backend/models/operations.py b/pyrit/backend/models/operations.py
index 214264e6e9..1b62d10bc0 100644
--- a/pyrit/backend/models/operations.py
+++ b/pyrit/backend/models/operations.py
@@ -7,8 +7,17 @@
from pydantic import BaseModel, ConfigDict, Field
from pyrit.backend.models.attacks import ScoreView
-from pyrit.models import AttackOutcome, Finding, FindingCreate, FindingEvidence, FindingSeverity, Operation
+from pyrit.models import (
+ AttackOutcome,
+ Finding,
+ FindingCreate,
+ FindingEvidence,
+ FindingSeverity,
+ Operation,
+ OperationUsage,
+)
from pyrit.models.harm_category import HarmCategory
+from pyrit.models.operation import OperationName
class FindingOptionsResponse(BaseModel):
@@ -29,12 +38,19 @@ class OperationListResponse(BaseModel):
items: list[OperationListItem]
+class OperationDetail(Operation):
+ """An operation with what refers to it by ID. Only an operation with no usage can be deleted."""
+
+ usage: OperationUsage
+
+
class OperationUpdateRequest(BaseModel):
- """The editable fields of an operation. The name cannot be changed here."""
+ """The editable fields of an operation. Omitted fields are unchanged."""
model_config = ConfigDict(extra="forbid")
- description: str
+ name: OperationName | None = None
+ description: str | None = None
class FindingListItem(Finding):
diff --git a/pyrit/backend/routes/operations.py b/pyrit/backend/routes/operations.py
index 8283342875..b7719df90d 100644
--- a/pyrit/backend/routes/operations.py
+++ b/pyrit/backend/routes/operations.py
@@ -14,6 +14,7 @@
FindingEvidenceListResponse,
FindingListResponse,
FindingOptionsResponse,
+ OperationDetail,
OperationListResponse,
OperationUpdateRequest,
)
@@ -25,6 +26,7 @@
OperationNotFoundError,
OperationService,
)
+from pyrit.memory.memory_interface import OperationInUseError
from pyrit.models import Finding, FindingCreate, Operation, OperationCreate
from pyrit.models.harm_category import HarmCategory
@@ -47,10 +49,7 @@ async def create_operation_async(request: OperationCreate) -> Operation:
try:
return await OperationService().create_async(request)
except DuplicateOperationError as exc:
- raise HTTPException(
- status_code=status.HTTP_409_CONFLICT,
- detail={"message": str(exc), "operation": exc.existing.model_dump(mode="json")},
- ) from exc
+ raise _name_conflict(exc) from exc
@router.get("", response_model=OperationListResponse)
@@ -75,26 +74,54 @@ async def get_finding_options_async() -> FindingOptionsResponse:
return FindingOptionsResponse(harm_types=list(HarmCategory))
-@router.get("/{operation_id}", response_model=Operation)
-async def get_operation_async(operation_id: UUID) -> Operation:
+@router.get("/{operation_id}", response_model=OperationDetail)
+async def get_operation_async(operation_id: UUID) -> OperationDetail:
"""
- Read one operation.
+ Read one operation with the counts of findings, attacks, and scenario runs that refer to it by ID.
Returns:
- Operation: The requested operation.
+ OperationDetail: The requested operation and its usage.
"""
- return await _require_operation_async(OperationService().get_async(operation_id))
+ return await _require_operation_async(OperationService().get_detail_async(operation_id))
@router.patch("/{operation_id}", response_model=Operation)
async def update_operation_async(*, operation_id: UUID, request: OperationUpdateRequest) -> Operation:
"""
- Replace an operation's description. The last save wins.
+ Rename an operation or replace its description. Omitted fields are unchanged and the last save wins.
Returns:
Operation: The updated operation.
+
+ Raises:
+ HTTPException: 409 with the existing operation when another operation has the name.
+ """
+ try:
+ return await _require_operation_async(
+ OperationService().update_async(operation_id=operation_id, request=request)
+ )
+ except DuplicateOperationError as exc:
+ raise _name_conflict(exc) from exc
+
+
+@router.delete("/{operation_id}", status_code=status.HTTP_204_NO_CONTENT)
+async def delete_operation_async(operation_id: UUID) -> Response:
+ """
+ Permanently remove an operation that no findings, attacks, or scenario runs refer to by ID.
+
+ Returns:
+ Response: An empty success response.
+
+ Raises:
+ HTTPException: 409 with the usage counts when the operation is in use.
"""
- return await _require_operation_async(OperationService().update_async(operation_id=operation_id, request=request))
+ try:
+ await _require_operation_async(OperationService().delete_async(operation_id))
+ except OperationInUseError as exc:
+ raise HTTPException(
+ status_code=status.HTTP_409_CONFLICT, detail={"message": str(exc), "usage": exc.usage.model_dump()}
+ ) from exc
+ return Response(status_code=status.HTTP_204_NO_CONTENT)
@router.post("/{operation_id}/findings", response_model=Finding, status_code=status.HTTP_201_CREATED)
@@ -225,3 +252,10 @@ async def _require_operation_async(pending: Awaitable[T]) -> T:
return await pending
except (OperationNotFoundError, FindingNotFoundError, FindingEvidenceNotFoundError) as exc:
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(exc)) from exc
+
+
+def _name_conflict(exc: DuplicateOperationError) -> HTTPException:
+ return HTTPException(
+ status_code=status.HTTP_409_CONFLICT,
+ detail={"message": str(exc), "operation": exc.existing.model_dump(mode="json")},
+ )
diff --git a/pyrit/backend/services/operation_service.py b/pyrit/backend/services/operation_service.py
index cfb133425a..f1d638b3ee 100644
--- a/pyrit/backend/services/operation_service.py
+++ b/pyrit/backend/services/operation_service.py
@@ -14,6 +14,7 @@
FindingEvidenceListResponse,
FindingListItem,
FindingListResponse,
+ OperationDetail,
OperationListItem,
OperationUpdateRequest,
)
@@ -96,22 +97,50 @@ async def get_async(self, operation_id: UUID) -> Operation:
except LookupError as exc:
raise OperationNotFoundError(f"Operation '{operation_id}' not found.") from exc
+ async def get_detail_async(self, operation_id: UUID) -> OperationDetail:
+ """
+ Read one operation with what refers to it by ID.
+
+ Returns:
+ OperationDetail: The operation and its usage.
+ """
+ operation = await self.get_async(operation_id)
+ usage = await CentralMemory.get_memory_instance().get_operation_usage_async(operation_id=operation_id)
+ return OperationDetail(**operation.model_dump(), usage=usage)
+
async def update_async(self, *, operation_id: UUID, request: OperationUpdateRequest) -> Operation:
"""
- Replace an operation's description.
+ Rename an operation or replace its description.
Returns:
Operation: The updated operation.
Raises:
OperationNotFoundError: If no operation has this ID.
+ DuplicateOperationError: If another operation already has the normalized name.
"""
try:
- return await CentralMemory.get_memory_instance().update_operation_async(
- operation_id=operation_id, description=request.description
+ saved = await CentralMemory.get_memory_instance().update_operation_async(
+ operation_id=operation_id, name=request.name, description=request.description
)
except LookupError as exc:
raise OperationNotFoundError(f"Operation '{operation_id}' not found.") from exc
+ if saved.id != operation_id:
+ raise DuplicateOperationError(saved)
+ return saved
+
+ async def delete_async(self, operation_id: UUID) -> None:
+ """
+ Permanently remove an operation that no findings, attacks, or scenario runs refer to by ID.
+
+ Raises:
+ OperationNotFoundError: If no operation has this ID.
+ OperationInUseError: If findings, attacks, or scenario runs refer to the operation by ID.
+ """
+ try:
+ await CentralMemory.get_memory_instance().delete_operation_async(operation_id=operation_id)
+ except LookupError as exc:
+ raise OperationNotFoundError(f"Operation '{operation_id}' not found.") from exc
async def create_finding_async(self, *, operation_id: UUID, request: FindingCreateRequest) -> Finding:
"""
@@ -131,7 +160,10 @@ async def create_finding_async(self, *, operation_id: UUID, request: FindingCrea
)
evidence = await self._new_evidence_async(finding_id=finding.id, request=request)
memory = CentralMemory.get_memory_instance()
- saved = await memory.add_finding_async(finding=finding, evidence=evidence)
+ try:
+ saved = await memory.add_finding_async(finding=finding, evidence=evidence)
+ except LookupError as exc:
+ raise OperationNotFoundError(f"Operation '{operation_id}' not found.") from exc
if saved.model_dump(exclude={"created_at"}) != finding.model_dump(exclude={"created_at"}) or (
evidence is not None
and await memory.get_finding_evidence_by_source_async(
diff --git a/pyrit/executor/attack/core/attack_executor.py b/pyrit/executor/attack/core/attack_executor.py
index 998befa66f..60b62d3175 100644
--- a/pyrit/executor/attack/core/attack_executor.py
+++ b/pyrit/executor/attack/core/attack_executor.py
@@ -466,7 +466,12 @@ async def run_one_async(index: int, params: AttackParameters) -> AttackStrategyR
persisted_result_ids[index] = context._persisted_attack_result_id
tasks = [run_one_async(i, p) for i, p in enumerate(params_list)]
- results_or_exceptions = await asyncio.gather(*tasks, return_exceptions=True)
+ # Admitting the batch's Operations once lets every objective finish if one is deleted midway.
+ operation_ids = {
+ operation_id for p in params_list if (operation_id := getattr(p, "operation_id", None)) is not None
+ }
+ async with attack._admit_saved_operations_async(operation_ids=operation_ids):
+ results_or_exceptions = await asyncio.gather(*tasks, return_exceptions=True)
return self._process_execution_results(
objectives=[p.objective for p in params_list],
diff --git a/pyrit/executor/attack/core/attack_strategy.py b/pyrit/executor/attack/core/attack_strategy.py
index cc02417b62..62ae6588e4 100644
--- a/pyrit/executor/attack/core/attack_strategy.py
+++ b/pyrit/executor/attack/core/attack_strategy.py
@@ -10,7 +10,8 @@
import traceback
import uuid
from abc import ABC
-from contextlib import asynccontextmanager
+from contextlib import asynccontextmanager, contextmanager
+from contextvars import ContextVar
from dataclasses import dataclass, field
from enum import Enum
from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, overload
@@ -54,7 +55,7 @@
from pyrit.score.observation.lifecycle import AttemptObserver
if TYPE_CHECKING:
- from collections.abc import AsyncGenerator, Callable
+ from collections.abc import AsyncGenerator, Callable, Collection, Generator
from types import TracebackType
from pyrit.executor.attack.component.prepended_conversation_config import (
@@ -76,6 +77,32 @@
AttackStrategyContextT = TypeVar("AttackStrategyContextT", bound="AttackContext[Any]")
AttackStrategyResultT = TypeVar("AttackStrategyResultT", bound="AttackResult")
+#: Operations that an enclosing execution or batch already admitted.
+_admitted_operation_ids: ContextVar[frozenset[uuid.UUID]] = ContextVar(
+ "pyrit_admitted_operation_ids", default=frozenset()
+)
+
+
+@contextmanager
+def _operation_admission_scope(operation_ids: Collection[uuid.UUID]) -> Generator[None, None, None]:
+ """
+ Treat these Operations as admitted for executions started within this scope.
+
+ Executions given an admitted Operation skip the existence check, so deleting it mid-run lets them finish
+ and save their results without it. Executions given any other Operation are still checked.
+
+ Args:
+ operation_ids (Collection[uuid.UUID]): The Operations to admit.
+
+ Yields:
+ None: Control returns to the caller with the Operations admitted.
+ """
+ token = _admitted_operation_ids.set(_admitted_operation_ids.get() | frozenset(operation_ids))
+ try:
+ yield
+ finally:
+ _admitted_operation_ids.reset(token)
+
def attack_outcome_from_score(score: Score) -> AttackOutcome:
"""
@@ -451,7 +478,31 @@ async def _persist_result_async(self, *, result: AttackStrategyResultT) -> None:
Args:
result (AttackStrategyResultT): The completed result to persist.
"""
- (await self._memory.add_attack_results_to_memory_async(attack_results=[result]))
+ await self._save_result_async(result=result)
+
+ async def _save_result_async(self, *, result: AttackResult) -> None:
+ """
+ Save an attack result, without its Operation if that Operation was deleted while the attack ran.
+
+ Args:
+ result (AttackResult): The result to save. Its Operation ID is cleared when the Operation is gone.
+
+ Raises:
+ LookupError: If the save fails for another missing reference.
+ """
+ try:
+ await self._memory.add_attack_results_to_memory_async(attack_results=[result])
+ except LookupError:
+ operation_id = result.operation_id
+ if operation_id is None or await self._memory.get_operations_async(operation_ids=[operation_id]):
+ raise
+ self._logger.warning(
+ "Operation %s was deleted while attack result %s ran; saving the result without an Operation.",
+ operation_id,
+ result.attack_result_id,
+ )
+ result.operation_id = None
+ await self._memory.add_attack_results_to_memory_async(attack_results=[result])
async def _confirm_result_persistence_async(
self, *, context: AttackStrategyContextT, attack_result_id: str
@@ -634,7 +685,7 @@ async def _on_error_async(
self._apply_targeted_harm_categories(context=context, result=error_result)
try:
- (await self._memory.add_attack_results_to_memory_async(attack_results=[error_result]))
+ await self._save_result_async(result=error_result)
except Exception as persistence_error:
context._error_result_persistence_error = persistence_error
await self._confirm_result_persistence_async(
@@ -969,10 +1020,14 @@ async def execute_with_context_async(self, *, context: AttackStrategyContextT) -
logger=self._logger,
)
context._objective_target_conversation_lifecycle = lifecycle
+ admitted = {context.params.operation_id} if context.params.operation_id is not None else set()
try:
async with lifecycle:
try:
- with attack_result_id_scope(attack_result_id=context._attack_result_id):
+ with (
+ attack_result_id_scope(attack_result_id=context._attack_result_id),
+ _operation_admission_scope(admitted),
+ ):
result = await super().execute_with_context_async(context=context)
except Exception as attack_error:
persistence_error = context._error_result_persistence_error
@@ -1004,7 +1059,8 @@ async def _resolve_operation_async(self, *, context: AttackStrategyContextT) ->
An explicit ID wins, and a ``memory_labels["operation"]`` name is then removed. Otherwise the name,
or one from ``GLOBAL_MEMORY_LABELS``, stays in the labels so prompts and the result keep it, and
links the execution to the saved Operation with that name when one exists. Executions attributed
- by a parent, such as a scenario, keep the labels and ID the parent resolved.
+ by a parent, such as a scenario, keep the labels and ID the parent resolved. An ID that an enclosing
+ execution or batch already admitted is not checked again, so a mid-run deletion doesn't stop the run.
Raises:
ValueError: If the ID is malformed.
@@ -1023,12 +1079,41 @@ async def _resolve_operation_async(self, *, context: AttackStrategyContextT) ->
elif operation_name is None and self._global_operation_name is not None:
operation_name = self._global_operation_name
context.memory_labels = {**labels, "operation": operation_name}
+ if operation_id is not None and operation_id in _admitted_operation_ids.get():
+ return
resolved = await CentralMemory.get_memory_instance().resolve_operation_id_async(
operation_id=operation_id, operation_name=operation_name
)
if resolved != operation_id:
context.params = dataclasses.replace(context.params, operation_id=resolved)
+ @asynccontextmanager
+ async def _admit_saved_operations_async(
+ self, *, operation_ids: Collection[uuid.UUID]
+ ) -> AsyncGenerator[None, None]:
+ """
+ Admit the given Operations that exist now for every execution started within this scope.
+
+ A batch admitted this way runs to completion if one of its Operations is deleted midway, and saves the
+ remaining results without it. Operations that don't exist when the scope opens are not admitted, so
+ executions given them are still rejected.
+
+ Args:
+ operation_ids (Collection[uuid.UUID]): The Operations the batch's executions refer to.
+
+ Yields:
+ None: Control returns to the caller with the saved Operations admitted.
+ """
+ pending = set(operation_ids) - _admitted_operation_ids.get()
+ memory = CentralMemory.get_memory_instance()
+ saved = (
+ {operation.id for operation in await memory.get_operations_async(operation_ids=list(pending))}
+ if pending
+ else set()
+ )
+ with _operation_admission_scope(saved):
+ yield
+
@asynccontextmanager
async def _execution_context_async(self, context: AttackStrategyContextT) -> AsyncGenerator[None, None]:
"""
diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py
index b522f82562..c4d971bd3b 100644
--- a/pyrit/memory/memory_interface.py
+++ b/pyrit/memory/memory_interface.py
@@ -48,6 +48,7 @@
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession
from sqlalchemy.orm import joinedload
from sqlalchemy.orm.attributes import InstrumentedAttribute, flag_modified
+from sqlalchemy.orm.exc import StaleDataError
from sqlalchemy.orm.session import Session
from sqlalchemy.sql import visitors
from sqlalchemy.sql.elements import ClauseElement
@@ -123,6 +124,7 @@
MessageScorable,
Observation,
Operation,
+ OperationUsage,
PromptDataType,
RetryEvent,
ScenarioAttackResultDelta,
@@ -192,6 +194,15 @@ class AttackStateConflictError(ValueError):
"""An atomic attack write no longer matches the state read by its caller."""
+class OperationInUseError(ValueError):
+ """An Operation still has findings or activity attributed to it by ID, so it cannot be deleted."""
+
+ def __init__(self, *, operation_id: uuid.UUID, usage: OperationUsage) -> None:
+ """Keep the usage so callers can say what still refers to the Operation."""
+ super().__init__(f"Operation '{operation_id}' is in use and cannot be deleted.")
+ self.usage = usage
+
+
#: Canonical criteria key of a seed that carries no conditions.
_NO_CONDITIONS_KEY = json.dumps([], separators=(",", ":"))
@@ -256,6 +267,20 @@ def _check_operation_name_change(*, existing: AttackResultEntry | ScenarioResult
)
+def _operation_ids(items: Sequence[object]) -> set[uuid.UUID]:
+ """
+ Collect the Operation IDs that attacks or scenario runs refer to, ignoring anything else.
+
+ Returns:
+ set[uuid.UUID]: The referenced Operation IDs.
+ """
+ return {
+ item.operation_id
+ for item in items
+ if isinstance(item, (AttackResult, AttackResultEntry, ScenarioResultEntry)) and item.operation_id is not None
+ }
+
+
@dataclass(frozen=True, slots=True, kw_only=True)
class _PreparedScorableContent:
"""A loose-content value prepared for durable database persistence."""
@@ -716,23 +741,98 @@ async def require_operation_async(self, *, operation_id: uuid.UUID) -> Operation
raise LookupError(f"Operation {operation_id} not found")
return operations[0]
- async def update_operation_async(self, *, operation_id: uuid.UUID, description: str) -> Operation:
+ async def update_operation_async(
+ self, *, operation_id: uuid.UUID, name: str | None = None, description: str | None = None
+ ) -> Operation:
"""
- Replace an operation's description. The last save wins; the name and identity are unchanged.
+ Rename an operation or replace its description. Omitted fields are unchanged and the last save wins.
+
+ Names stay unique, ignoring case and surrounding whitespace, so a rename that only changes case is allowed.
Returns:
- Operation: The updated operation.
+ Operation: The updated operation, or the other operation that already has the name, unchanged.
Raises:
LookupError: If the Operation does not exist.
+ IntegrityError: If the update fails for a reason other than a taken name.
"""
- async with await self.get_session_async() as session, session.begin():
- entry = await session.get(OperationEntry, operation_id)
- if entry is None:
- raise LookupError(f"Operation {operation_id} not found")
- entry.description = description
- await session.flush()
- return entry.get_operation()
+ changes = {key: value for key, value in {"name": name, "description": description}.items() if value is not None}
+ try:
+ async with await self.get_session_async() as session, session.begin():
+ entry = await session.get(OperationEntry, operation_id)
+ if entry is None:
+ raise LookupError(f"Operation {operation_id} not found")
+ updated = Operation.model_validate(entry.get_operation().model_dump() | changes)
+ entry.name, entry.name_key = updated.name, operation_name_key(updated.name)
+ entry.description = updated.description
+ await session.flush()
+ return entry.get_operation()
+ except StaleDataError as exc:
+ # A concurrent delete landed between the read and the update.
+ raise LookupError(f"Operation {operation_id} not found") from exc
+ except IntegrityError:
+ existing = await self.get_operations_async(name=name) if name is not None else []
+ if not existing:
+ raise
+ return existing[0]
+
+ async def get_operation_usage_async(self, *, operation_id: uuid.UUID) -> OperationUsage:
+ """
+ Count the findings, attacks, and scenario runs that refer to an Operation by ID, in one indexed query.
+
+ Returns:
+ OperationUsage: The counts, which are all zero for an unknown ID.
+ """
+ counts = [
+ select(func.count())
+ .select_from(table)
+ .where(table.operation_id == operation_id)
+ .scalar_subquery()
+ .label(field)
+ for field, table in (
+ ("findings", FindingEntry),
+ ("attacks", AttackResultEntry),
+ ("scenario_runs", ScenarioResultEntry),
+ )
+ ]
+ async with await self.get_session_async() as session:
+ return OperationUsage.model_validate(dict((await session.execute(select(*counts))).one()._mapping))
+
+ async def delete_operation_async(self, *, operation_id: uuid.UUID) -> None:
+ """
+ Permanently remove an Operation that no findings, attacks, or scenario runs refer to by ID.
+
+ Activity attributed only by a deprecated Operation name does not prevent deletion and is unchanged.
+
+ Raises:
+ LookupError: If the Operation does not exist.
+ OperationInUseError: If findings, attacks, or scenario runs refer to the Operation by ID.
+ RuntimeError: If the database does not return a cursor result.
+ """
+ # SQLite does not enforce foreign keys, so the delete statement itself checks for references.
+ unreferenced = [
+ ~select(table.operation_id).where(table.operation_id == operation_id).exists()
+ for table in (FindingEntry, AttackResultEntry, ScenarioResultEntry)
+ ]
+ try:
+ async with await self.get_session_async() as session, session.begin():
+ result = await session.execute(
+ delete(OperationEntry)
+ .where(OperationEntry.id == operation_id, *unreferenced)
+ .execution_options(synchronize_session=False)
+ )
+ if not isinstance(result, CursorResult):
+ raise RuntimeError("Operation deletion did not return a database cursor result.")
+ if result.rowcount > 0:
+ return
+ if await session.get(OperationEntry, operation_id) is None:
+ raise LookupError(f"Operation {operation_id} not found")
+ except IntegrityError as exc:
+ # SQL Server's foreign keys reject the delete when a reference commits while it runs.
+ usage = await self.get_operation_usage_async(operation_id=operation_id)
+ raise OperationInUseError(operation_id=operation_id, usage=usage) from exc
+ usage = await self.get_operation_usage_async(operation_id=operation_id)
+ raise OperationInUseError(operation_id=operation_id, usage=usage)
async def resolve_operation_id_async(
self, *, operation_id: uuid.UUID | str | None, operation_name: str | None = None
@@ -773,6 +873,7 @@ async def add_finding_async(self, *, finding: Finding, evidence: FindingEvidence
Finding: The saved finding, or the finding already saved with this ID.
Raises:
+ LookupError: If the finding's Operation does not exist.
IntegrityError: If the insert fails for a reason other than a taken ID.
"""
entry = FindingEntry(finding)
@@ -780,16 +881,22 @@ async def add_finding_async(self, *, finding: Finding, evidence: FindingEvidence
async with await self.get_session_async() as session, session.begin():
session.add(entry)
await session.flush()
+ # Checked after the insert so the write lock is held; SQLite does not enforce the foreign key.
+ if await session.get(OperationEntry, finding.operation_id) is None:
+ raise LookupError(f"Operation {finding.operation_id} not found")
if evidence is not None:
session.add(FindingEvidenceEntry(evidence))
await session.flush()
return entry.get_finding()
- except IntegrityError:
+ except IntegrityError as exc:
async with await self.get_session_async() as session:
existing = await session.get(FindingEntry, finding.id)
- if existing is None:
- raise
- return existing.get_finding()
+ if existing is not None:
+ return existing.get_finding()
+ # SQL Server's foreign key rejects the insert when the Operation was deleted meanwhile.
+ if await session.get(OperationEntry, finding.operation_id) is None:
+ raise LookupError(f"Operation {finding.operation_id} not found") from exc
+ raise
async def update_finding_async(
self, *, operation_id: uuid.UUID, finding_id: uuid.UUID, request: FindingCreate
@@ -2367,13 +2474,18 @@ def _update_entry(self, entry: Base) -> None:
"""
with closing(self._get_session()) as session:
try:
+ existing = None
if isinstance(entry, (AttackResultEntry, ScenarioResultEntry)):
existing = session.get(type(entry), entry.id)
self._validate_activity_operation(session=session, entry=entry, existing=existing)
session.merge(entry)
+ if existing is None:
+ self._recheck_activity_operations(session=session, entries=[entry])
session.commit()
except SQLAlchemyError as e:
session.rollback()
+ if isinstance(e, IntegrityError):
+ self._raise_if_operations_missing(session=session, operation_ids=_operation_ids([entry]), error=e)
logger.exception(f"Error updating entry in the table: {e}")
raise
@@ -5594,11 +5706,18 @@ def _execute_add_attack_results_to_memory(self, *, attack_results: Sequence[Atta
"""
with closing(self._get_session()) as session:
try:
- for attack_result in attack_results:
+ entries = [
self._add_attack_result_to_session(session=session, attack_result=attack_result)
+ for attack_result in attack_results
+ ]
+ self._recheck_activity_operations(session=session, entries=entries)
session.commit()
- except SQLAlchemyError:
+ except SQLAlchemyError as error:
session.rollback()
+ if isinstance(error, IntegrityError):
+ self._raise_if_operations_missing(
+ session=session, operation_ids=_operation_ids(attack_results), error=error
+ )
raise
def _add_attack_result_to_session(self, *, session: Session, attack_result: AttackResult) -> AttackResultEntry:
@@ -5647,6 +5766,37 @@ def _validate_activity_operation(
if entry.operation_id is not None and session.get(OperationEntry, entry.operation_id) is None:
raise LookupError(f"Operation {entry.operation_id} not found")
+ @staticmethod
+ def _recheck_activity_operations(*, session: Session, entries: Sequence[Base]) -> None:
+ """
+ Check new activity's Operations again after inserting it, while this transaction holds the write lock.
+
+ Raises:
+ LookupError: If an Operation was deleted after the first check.
+ """
+ # SQLite does not enforce foreign keys, so a concurrent delete would otherwise orphan the activity.
+ session.flush()
+ MemoryInterface._raise_if_operations_missing(session=session, operation_ids=_operation_ids(entries))
+
+ @staticmethod
+ def _raise_if_operations_missing(
+ *, session: Session, operation_ids: set[uuid.UUID], error: Exception | None = None
+ ) -> None:
+ """
+ Report activity whose Operation no longer exists as a missing Operation.
+
+ After a rolled-back insert, this separates a foreign-key failure from a deleted Operation, which
+ SQL Server raises before the post-insert check runs, from unrelated integrity errors.
+
+ Raises:
+ LookupError: If one of the Operations does not exist.
+ """
+ if not operation_ids:
+ return
+ saved = set(session.scalars(select(OperationEntry.id).where(OperationEntry.id.in_(operation_ids))))
+ if missing := operation_ids - saved:
+ raise LookupError(f"Operation {next(iter(missing))} not found") from error
+
def _execute_add_conversation_branches_to_attack(
self,
*,
@@ -5769,8 +5919,10 @@ def _execute_add_conversation_branches_to_attack(
f"{receipt_key}:{request_fingerprint}": True,
}
entry.timestamp = datetime.now(UTC)
+ if new_attack is not None:
+ self._recheck_activity_operations(session=session, entries=[entry])
return True
- except IntegrityError:
+ except IntegrityError as error:
# A concurrent retry may have committed the same new attack before our insert.
with closing(self._get_session()) as session:
existing = session.get(AttackResultEntry, uuid.UUID(attack_result_id))
@@ -5783,6 +5935,9 @@ def _execute_add_conversation_branches_to_attack(
)
):
return False
+ self._raise_if_operations_missing(
+ session=session, operation_ids=_operation_ids([new_attack]), error=error
+ )
raise
def _execute_update_attack_result_conditionally(
@@ -6652,9 +6807,14 @@ def _execute_add_scenario_results_to_memory(self, *, scenario_results: Sequence[
for entry in entries:
self._validate_activity_operation(session=session, entry=entry)
session.add_all(entries)
+ self._recheck_activity_operations(session=session, entries=entries)
session.commit()
- except SQLAlchemyError:
+ except SQLAlchemyError as error:
session.rollback()
+ if isinstance(error, IntegrityError):
+ self._raise_if_operations_missing(
+ session=session, operation_ids=_operation_ids(entries), error=error
+ )
raise
@classmethod
diff --git a/pyrit/models/__init__.py b/pyrit/models/__init__.py
index 1fcd1483bd..2851702be8 100644
--- a/pyrit/models/__init__.py
+++ b/pyrit/models/__init__.py
@@ -128,7 +128,7 @@
)
from pyrit.models.messages.conversation_reference import ConversationReference, ConversationType
from pyrit.models.messages.conversation_retry import ConversationRetry, ConversationRetryReason
- from pyrit.models.operation import Operation, OperationCreate
+ from pyrit.models.operation import Operation, OperationCreate, OperationUsage
from pyrit.models.parameter import (
ComponentType,
Parameter,
@@ -406,6 +406,7 @@
"ObjectiveTargetEvaluationIdentifier": "pyrit.models.identifiers",
"Operation": "pyrit.models.operation",
"OperationCreate": "pyrit.models.operation",
+ "OperationUsage": "pyrit.models.operation",
"Parameter": "pyrit.models.parameter",
"ParameterDestination": "pyrit.models.parameter",
"PresetEntry": "pyrit.models.catalog",
diff --git a/pyrit/models/operation.py b/pyrit/models/operation.py
index c4bbbe9d51..a8d0052de4 100644
--- a/pyrit/models/operation.py
+++ b/pyrit/models/operation.py
@@ -2,9 +2,10 @@
# Licensed under the MIT license.
from datetime import UTC, datetime
+from typing import Annotated
from uuid import UUID, uuid4
-from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, field_validator
+from pydantic import AfterValidator, AwareDatetime, BaseModel, ConfigDict, Field
def operation_name_key(name: str) -> str:
@@ -17,36 +18,46 @@ def operation_name_key(name: str) -> str:
return name.strip().casefold()
+def _normalize_operation_name(value: str) -> str:
+ """
+ Trim surrounding whitespace and enforce the SQL Server NVARCHAR(128) bound.
+
+ Returns:
+ str: The trimmed display name.
+
+ Raises:
+ ValueError: If the trimmed name is blank or exceeds 128 UTF-16 code units.
+ """
+ value = value.strip()
+ if not value:
+ raise ValueError("Must not be blank.")
+ if len(value.encode("utf-16-le")) // 2 > 128:
+ raise ValueError("Must be at most 128 UTF-16 code units.")
+ return value
+
+
+OperationName = Annotated[str, AfterValidator(_normalize_operation_name)]
+
+
class OperationCreate(BaseModel):
"""A named engagement that groups human findings."""
model_config = ConfigDict(extra="forbid")
- name: str
+ name: OperationName
description: str = ""
- @field_validator("name")
- @classmethod
- def _normalize_name(cls, value: str) -> str:
- """
- Trim surrounding whitespace and enforce the SQL Server NVARCHAR(128) bound.
-
- Returns:
- str: The trimmed display name.
-
- Raises:
- ValueError: If the trimmed name is blank or exceeds 128 UTF-16 code units.
- """
- value = value.strip()
- if not value:
- raise ValueError("Must not be blank.")
- if len(value.encode("utf-16-le")) // 2 > 128:
- raise ValueError("Must be at most 128 UTF-16 code units.")
- return value
-
class Operation(OperationCreate):
"""A saved operation with stable identity and creation time."""
id: UUID = Field(default_factory=uuid4)
created_at: AwareDatetime = Field(default_factory=lambda: datetime.now(UTC))
+
+
+class OperationUsage(BaseModel):
+ """What still refers to an Operation by ID. Activity attributed only by a deprecated name is not counted."""
+
+ findings: int = Field(ge=0)
+ attacks: int = Field(ge=0)
+ scenario_runs: int = Field(ge=0)
diff --git a/tests/unit/backend/test_operations.py b/tests/unit/backend/test_operations.py
index 754ec873b1..c633211fde 100644
--- a/tests/unit/backend/test_operations.py
+++ b/tests/unit/backend/test_operations.py
@@ -34,6 +34,7 @@
pytestmark = pytest.mark.usefixtures("patch_central_database")
MISSING_ID = UUID(int=404)
+EMPTY_USAGE = {"findings": 0, "attacks": 0, "scenario_runs": 0}
def test_operations_create_list_and_get(compatibility_headers: dict[str, str]) -> None:
@@ -45,13 +46,11 @@ def test_operations_create_list_and_get(compatibility_headers: dict[str, str]) -
client.post("/api/operations", json={"name": "Alpha"})
assert [item["name"] for item in client.get("/api/operations").json()["items"]] == ["Alpha", "Red team / α%"]
assert all(item["finding_counts"] == {} for item in client.get("/api/operations").json()["items"])
- assert client.get(f"/api/operations/{operation['id']}").json() == operation
+ assert client.get(f"/api/operations/{operation['id']}").json() == {**operation, "usage": EMPTY_USAGE}
assert client.get(f"/api/operations/{MISSING_ID}").status_code == 404
-def test_operation_description_is_set_on_create_and_edited_without_renaming(
- compatibility_headers: dict[str, str],
-) -> None:
+def test_operation_description_is_set_on_create_and_edited(compatibility_headers: dict[str, str]) -> None:
client = TestClient(app, headers=compatibility_headers)
operation = client.post("/api/operations", json={"name": "Scoped", "description": "Web app only"}).json()
assert operation["description"] == "Web app only"
@@ -59,14 +58,67 @@ def test_operation_description_is_set_on_create_and_edited_without_renaming(
updated = client.patch(path, json={"description": "Web app and API\nNo social engineering"})
assert updated.status_code == 200
assert updated.json() == {**operation, "description": "Web app and API\nNo social engineering"}
- assert client.get(path).json() == updated.json()
+ assert client.get(path).json() == {**updated.json(), "usage": EMPTY_USAGE}
assert client.patch(path, json={"description": ""}).json()["description"] == ""
- assert client.patch(path, json={"description": "x", "name": "Renamed"}).status_code == 422
+ assert client.patch(path, json={"description": "x", "created_at": operation["created_at"]}).status_code == 422
assert client.patch(f"/api/operations/{MISSING_ID}", json={"description": "x"}).status_code == 404
assert TestClient(app).patch(path, json={"description": "x"}).status_code != 200
assert client.get("/api/operations").json()["items"][0]["description"] == ""
+def test_operation_rename_keeps_identity_and_rejects_taken_or_invalid_names(
+ compatibility_headers: dict[str, str],
+) -> None:
+ client = TestClient(app, headers=compatibility_headers)
+ operation = client.post("/api/operations", json={"name": "Scoped", "description": "Web app"}).json()
+ other = client.post("/api/operations", json={"name": "Other"}).json()
+ path = f"/api/operations/{operation['id']}"
+ recased = client.patch(path, json={"name": " SCOPED "})
+ assert recased.status_code == 200
+ assert recased.json() == {**operation, "name": "SCOPED"}
+ conflict = client.patch(path, json={"name": "other", "description": "Lost"})
+ assert conflict.status_code == 409
+ assert conflict.json()["detail"]["operation"] == other
+ for name in [" ", "😀" * 65]:
+ assert client.patch(path, json={"name": name}).status_code == 422
+ assert client.get(path).json() == {**recased.json(), "usage": EMPTY_USAGE}
+ assert client.patch(path, json={"name": "Engagement / α%"}).json() == {**operation, "name": "Engagement / α%"}
+ assert [item["name"] for item in client.get("/api/operations").json()["items"]] == ["Engagement / α%", "Other"]
+
+
+def test_empty_operation_is_deleted(compatibility_headers: dict[str, str]) -> None:
+ client = TestClient(app, headers=compatibility_headers)
+ operation = client.post("/api/operations", json={"name": "Typo"}).json()
+ path = f"/api/operations/{operation['id']}"
+ assert TestClient(app).delete(path).status_code != 204
+ deleted = client.delete(path)
+ assert deleted.status_code == 204
+ assert deleted.content == b""
+ assert client.get(path).status_code == 404
+ assert client.delete(path).status_code == 404
+ assert client.patch(path, json={"name": "Fixed"}).status_code == 404
+ assert client.get("/api/operations").json()["items"] == []
+
+
+async def test_finding_for_an_operation_deleted_during_the_request_is_not_found_async(
+ sqlite_instance: SQLiteMemory, compatibility_headers: dict[str, str]
+) -> None:
+ operation = await sqlite_instance.add_operation_async(Operation(name="Case"))
+ original_get_async = OperationService.get_async
+
+ async def get_then_delete_async(self: OperationService, operation_id: UUID) -> Operation:
+ found = await original_get_async(self, operation_id)
+ await sqlite_instance.delete_operation_async(operation_id=operation_id)
+ return found
+
+ with patch.object(OperationService, "get_async", get_then_delete_async):
+ response = TestClient(app, headers=compatibility_headers).post(
+ f"/api/operations/{operation.id}/findings", json={"title": "Late", "severity": "low"}
+ )
+ assert response.status_code == 404
+ assert await sqlite_instance.get_findings_async() == []
+
+
def test_duplicate_operation_conflicts_with_existing_record(compatibility_headers: dict[str, str]) -> None:
client = TestClient(app, headers=compatibility_headers)
existing = client.post("/api/operations", json={"name": "Operation A"}).json()
@@ -252,7 +304,7 @@ def test_finding_update_and_delete_preserve_identity_and_operation(
assert deleted.status_code == 204
assert deleted.content == b""
assert client.get(path).json()["items"] == []
- assert client.get(f"/api/operations/{operation['id']}").json() == operation
+ assert client.get(f"/api/operations/{operation['id']}").json() == {**operation, "usage": EMPTY_USAGE}
assert client.delete(finding_path).status_code == 404
assert client.put(finding_path, json={"title": "Missing", "severity": "low"}).status_code == 404
@@ -332,6 +384,24 @@ def test_finding_evidence_attach_returns_created_then_already_attached(
assert client.delete(f"{path}/{evidence_id}").status_code == 404
+def test_operation_in_use_is_not_deleted_and_id_linked_activity_follows_a_rename(
+ compatibility_headers: dict[str, str], evidence_source_async: tuple[Operation, Finding, AttackResult]
+) -> None:
+ operation, finding, attack = evidence_source_async
+ client = TestClient(app, headers=compatibility_headers)
+ path = f"/api/operations/{operation.id}"
+ usage = {"findings": 1, "attacks": 1, "scenario_runs": 0}
+ assert client.get(path).json()["usage"] == usage
+ refused = client.delete(path)
+ assert refused.status_code == 409
+ assert refused.json()["detail"]["usage"] == usage
+ assert client.patch(path, json={"name": "Renamed"}).status_code == 200
+ evidence_path = f"{path}/findings/{finding.id}/evidence"
+ client.post(evidence_path, json={"attack_result_id": attack.attack_result_id})
+ assert client.get(evidence_path).json()["items"][0]["source_operation_name"] == "Renamed"
+ assert client.get(path).json()["name"] == "Renamed"
+
+
@pytest.mark.parametrize(
("attribution", "source_name"),
[("unassigned", None), ("other_operation", "Other"), ("legacy_exact_name", "Exact Operation")],
diff --git a/tests/unit/executor/attack/core/test_operation_id_attribution.py b/tests/unit/executor/attack/core/test_operation_id_attribution.py
index efe464a25e..41ae61f01a 100644
--- a/tests/unit/executor/attack/core/test_operation_id_attribution.py
+++ b/tests/unit/executor/attack/core/test_operation_id_attribution.py
@@ -1,19 +1,29 @@
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
-from unittest.mock import patch
+import asyncio
+import logging
+from unittest.mock import AsyncMock, patch
from uuid import uuid4
import pytest
from unit.mocks import MockPromptTarget
-from pyrit.executor.attack import PromptSendingAttack
-from pyrit.executor.attack.compound import SequentialAttack, SequentialChildAttack
+from pyrit.executor.attack import AttackExecutor, PromptSendingAttack
+from pyrit.executor.attack.compound import SequenceCompletionPolicy, SequentialAttack, SequentialChildAttack
from pyrit.executor.attack.core.attack_parameters import AttackParameters
from pyrit.executor.attack.core.attack_result_attribution import AttackResultAttribution
from pyrit.executor.attack.core.attack_strategy import AttackContext, AttackStrategy
-from pyrit.memory import MemoryInterface
-from pyrit.models import AttackResult, AttackResultSelection, AttackSeedGroup, Operation, SeedObjective
+from pyrit.memory import CentralMemory, MemoryInterface
+from pyrit.models import (
+ AttackResult,
+ AttackResultSelection,
+ AttackSeedGroup,
+ Message,
+ Operation,
+ OperationUsage,
+ SeedObjective,
+)
class AdmissionProbe(AttackStrategy[AttackContext[AttackParameters], AttackResult]):
@@ -127,3 +137,138 @@ async def test_compound_children_inherit_id_async(self, sqlite_instance: MemoryI
saved = await sqlite_instance.get_attack_results_async(result_selection=AttackResultSelection.ALL_RESULTS)
assert len(saved) == 2
assert all(row.operation_id == operation.id for row in saved)
+
+
+class _PausedTarget(MockPromptTarget):
+ """A target that waits to respond until the test releases it."""
+
+ def __init__(self) -> None:
+ super().__init__()
+ self.sending = asyncio.Event()
+ self.release = asyncio.Event()
+
+ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]:
+ self.sending.set()
+ await self.release.wait()
+ return await super()._send_prompt_to_target_async(normalized_conversation=normalized_conversation)
+
+
+class _DeletingProbe(AdmissionProbe):
+ """Deletes the run's Operation from inside the attack, then succeeds or fails."""
+
+ async def _perform_async(self, *, context: AttackContext[AttackParameters]) -> AttackResult:
+ operation_id = context.params.operation_id
+ assert operation_id is not None
+ await CentralMemory.get_memory_instance().delete_operation_async(operation_id=operation_id)
+ return await super()._perform_async(context=context)
+
+
+@pytest.mark.usefixtures("patch_central_database")
+class TestOperationDeletedDuringRun:
+ async def test_in_flight_attack_keeps_its_result_without_the_deleted_operation_async(
+ self, sqlite_instance: MemoryInterface, caplog: pytest.LogCaptureFixture
+ ) -> None:
+ operation = await sqlite_instance.add_operation_async(Operation(name="Deleted mid-run"))
+ target = _PausedTarget()
+ running = asyncio.create_task(
+ PromptSendingAttack(objective_target=target).execute_async(
+ objective="Say hello.", operation_id=operation.id
+ )
+ )
+ await asyncio.wait_for(target.sending.wait(), timeout=10)
+ # Nothing refers to the Operation until the result is saved, so it can be deleted mid-run.
+ assert await sqlite_instance.get_operation_usage_async(operation_id=operation.id) == OperationUsage(
+ findings=0, attacks=0, scenario_runs=0
+ )
+ await sqlite_instance.delete_operation_async(operation_id=operation.id)
+ target.release.set()
+ with caplog.at_level(logging.WARNING):
+ result = await asyncio.wait_for(running, timeout=10)
+ assert result.operation_id is None
+ saved = await sqlite_instance.get_attack_results_async(attack_result_ids=[result.attack_result_id])
+ assert [row.operation_id for row in saved] == [None]
+ assert len(await sqlite_instance.get_message_pieces_async(conversation_id=result.conversation_id)) == 2
+ assert str(operation.id) in caplog.text
+
+ @pytest.mark.parametrize("fail", [False, True])
+ async def test_success_and_error_results_are_saved_without_the_deleted_operation_async(
+ self, sqlite_instance: MemoryInterface, fail: bool
+ ) -> None:
+ operation = await sqlite_instance.add_operation_async(Operation(name="Deleted mid-run"))
+ attack = _DeletingProbe(fail=fail)
+ if fail:
+ with pytest.raises(RuntimeError, match="Probe failed"):
+ await attack.execute_async(objective="objective", operation_id=operation.id)
+ else:
+ assert (await attack.execute_async(objective="objective", operation_id=operation.id)).operation_id is None
+ saved = await sqlite_instance.get_attack_results_async(result_selection=AttackResultSelection.ALL_RESULTS)
+ assert [row.operation_id for row in saved] == [None]
+ assert await sqlite_instance.get_operations_async() == []
+
+ async def test_other_lookup_errors_still_fail_the_save_async(self, sqlite_instance: MemoryInterface) -> None:
+ operation = await sqlite_instance.add_operation_async(Operation(name="Still here"))
+ save = AsyncMock(side_effect=LookupError("Atomic attack identifier not found"))
+ with patch.object(sqlite_instance, "add_attack_results_to_memory_async", save):
+ with pytest.raises(LookupError, match="identifier"):
+ await AdmissionProbe().execute_async(objective="objective", operation_id=operation.id)
+ save.assert_awaited_once()
+
+ async def test_in_flight_sequence_runs_every_child_without_the_deleted_operation_async(
+ self, sqlite_instance: MemoryInterface
+ ) -> None:
+ operation = await sqlite_instance.add_operation_async(Operation(name="Deleted mid-sequence"))
+ target = _PausedTarget()
+ children = [
+ SequentialChildAttack(
+ strategy=PromptSendingAttack(objective_target=target),
+ seed_group=AttackSeedGroup(seeds=[SeedObjective(value=f"Say hello {index}.")]),
+ )
+ for index in range(2)
+ ]
+ sequence = SequentialAttack(
+ objective_target=target, child_attacks=children, completion_policy=SequenceCompletionPolicy.EXHAUSTIVE
+ )
+ running = asyncio.create_task(sequence.execute_async(objective="Say hello.", operation_id=operation.id))
+ await asyncio.wait_for(target.sending.wait(), timeout=10)
+ await sqlite_instance.delete_operation_async(operation_id=operation.id)
+ target.release.set()
+ result = await asyncio.wait_for(running, timeout=10)
+ assert result.operation_id is None
+ assert len(target.prompt_sent) == 2
+ saved = await sqlite_instance.get_attack_results_async(result_selection=AttackResultSelection.ALL_RESULTS)
+ assert len(saved) == 3
+ assert all(row.operation_id is None for row in saved)
+
+ async def test_in_flight_batch_runs_every_objective_without_the_deleted_operation_async(
+ self, sqlite_instance: MemoryInterface
+ ) -> None:
+ operation = await sqlite_instance.add_operation_async(Operation(name="Deleted mid-batch"))
+ target = _PausedTarget()
+ running = asyncio.create_task(
+ AttackExecutor().execute_attack_async(
+ attack=PromptSendingAttack(objective_target=target),
+ objectives=["Say hello.", "Say goodbye."],
+ operation_id=operation.id,
+ )
+ )
+ await asyncio.wait_for(target.sending.wait(), timeout=10)
+ await sqlite_instance.delete_operation_async(operation_id=operation.id)
+ target.release.set()
+ executor_result = await asyncio.wait_for(running, timeout=10)
+ assert [result.operation_id for result in executor_result.completed_results] == [None, None]
+ assert len(target.prompt_sent) == 2
+ saved = await sqlite_instance.get_attack_results_async(result_selection=AttackResultSelection.ALL_RESULTS)
+ assert [row.operation_id for row in saved] == [None, None]
+
+ async def test_batch_given_a_missing_operation_still_rejects_every_objective_async(
+ self, sqlite_instance: MemoryInterface
+ ) -> None:
+ attack = AdmissionProbe()
+ executor_result = await AttackExecutor().execute_attack_async(
+ attack=attack, objectives=["first", "second"], operation_id=uuid4(), return_partial_on_failure=True
+ )
+ assert executor_result.completed_results == []
+ assert all(isinstance(error, LookupError) for _, error in executor_result.incomplete_objectives)
+ assert len(executor_result.incomplete_objectives) == 2
+ assert attack.setup_count == attack.perform_count == 0
+ assert await sqlite_instance.get_attack_results_async(result_selection=AttackResultSelection.ALL_RESULTS) == []
diff --git a/tests/unit/memory/test_activity_operation_ids.py b/tests/unit/memory/test_activity_operation_ids.py
index f9b722337b..7e42280d5e 100644
--- a/tests/unit/memory/test_activity_operation_ids.py
+++ b/tests/unit/memory/test_activity_operation_ids.py
@@ -1,17 +1,33 @@
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
+from collections.abc import AsyncGenerator, AsyncIterator
+from contextlib import asynccontextmanager
from pathlib import Path
+from typing import Any
+from unittest.mock import patch
from uuid import UUID, uuid4
import pytest
import sqlalchemy as sa
from alembic import command
from alembic.config import Config
+from sqlalchemy import event
+from sqlalchemy.exc import IntegrityError
+from sqlalchemy.ext.asyncio import AsyncSession
-from pyrit.memory import MemoryInterface
+from pyrit.memory import MemoryInterface, SQLiteMemory
+from pyrit.memory.memory_interface import OperationInUseError
from pyrit.memory.memory_models import AttackResultEntry, ScenarioResultEntry
-from pyrit.models import AttackResult, ComponentIdentifier, Operation, ScenarioResult
+from pyrit.models import (
+ AttackResult,
+ ComponentIdentifier,
+ Conversation,
+ Finding,
+ Operation,
+ OperationUsage,
+ ScenarioResult,
+)
from unit.mocks import make_scenario_result
REVISION = "c8d3e5f7a901"
@@ -295,3 +311,151 @@ def test_upgrade_registers_legacy_names_without_changing_activity() -> None:
command.downgrade(config, PARENT)
finally:
engine.dispose()
+
+
+async def _refer_to_operation_async(memory: MemoryInterface, *, operation_id: UUID, reference: str) -> None:
+ if reference == "finding":
+ await memory.add_finding_async(finding=Finding(operation_id=operation_id, title="Assessment", severity="low"))
+ elif reference == "attack":
+ await memory.add_attack_results_to_memory_async(
+ attack_results=[AttackResult(conversation_id="linked", objective="objective", operation_id=operation_id)]
+ )
+ else:
+ scenario = _scenario()
+ scenario.operation_id = operation_id
+ await memory.add_scenario_results_to_memory_async(scenario_results=[scenario])
+
+
+@pytest.mark.parametrize("reference", ["finding", "attack", "scenario"])
+async def test_delete_operation_refuses_any_id_reference_async(
+ sqlite_instance: MemoryInterface, reference: str
+) -> None:
+ operation = await sqlite_instance.add_operation_async(Operation(name="Engagement"))
+ await _refer_to_operation_async(sqlite_instance, operation_id=operation.id, reference=reference)
+ usage = OperationUsage(
+ findings=int(reference == "finding"),
+ attacks=int(reference == "attack"),
+ scenario_runs=int(reference == "scenario"),
+ )
+ assert await sqlite_instance.get_operation_usage_async(operation_id=operation.id) == usage
+ with pytest.raises(OperationInUseError) as raised:
+ await sqlite_instance.delete_operation_async(operation_id=operation.id)
+ assert raised.value.usage == usage
+ assert await sqlite_instance.require_operation_async(operation_id=operation.id) == operation
+
+
+async def test_delete_operation_keeps_name_only_activity_and_other_operations_async(
+ sqlite_instance: MemoryInterface,
+) -> None:
+ operation = await sqlite_instance.add_operation_async(Operation(name="Engagement"))
+ other = await sqlite_instance.add_operation_async(Operation(name="Other"))
+ await _refer_to_operation_async(sqlite_instance, operation_id=other.id, reference="attack")
+ legacy = AttackResult(conversation_id="legacy", objective="objective", operation=operation.name)
+ with pytest.warns(DeprecationWarning, match="operation_id"):
+ await sqlite_instance.add_attack_results_to_memory_async(attack_results=[legacy])
+ empty = OperationUsage(findings=0, attacks=0, scenario_runs=0)
+ assert await sqlite_instance.get_operation_usage_async(operation_id=operation.id) == empty
+ await sqlite_instance.delete_operation_async(operation_id=operation.id)
+ assert await sqlite_instance.get_operations_async() == [other]
+ found = await sqlite_instance.get_attack_results_async(operation=[operation.name])
+ assert [(result.attack_result_id, result.operation_id) for result in found] == [(legacy.attack_result_id, None)]
+ with pytest.raises(LookupError, match="Operation"):
+ await sqlite_instance.delete_operation_async(operation_id=operation.id)
+
+
+def _enforce_foreign_keys(dbapi_connection: Any, _: Any) -> None:
+ cursor = dbapi_connection.cursor()
+ cursor.execute("PRAGMA foreign_keys=ON")
+ cursor.close()
+
+
+@asynccontextmanager
+async def _foreign_key_memory_async(db_path: Path) -> AsyncGenerator[SQLiteMemory, None]:
+ """Open a file database that enforces foreign keys, as SQL Server does."""
+ # SQLiteMemory(...) would return the sqlite_instance singleton when that fixture is active.
+ memory = SQLiteMemory.__new__(SQLiteMemory)
+ with patch.object(memory, "cleanup"):
+ memory.__init__(db_path=db_path, silent=True)
+ for engine in (memory.engine, memory._get_async_engine().sync_engine):
+ event.listen(engine, "connect", _enforce_foreign_keys)
+ memory.engine.dispose()
+ try:
+ yield memory
+ finally:
+ await memory.dispose_engine_async()
+
+
+@pytest.fixture
+async def foreign_key_memory(tmp_path: Path) -> AsyncIterator[SQLiteMemory]:
+ async with _foreign_key_memory_async(tmp_path / "foreign-keys.db") as memory:
+ yield memory
+
+
+@pytest.fixture(params=[False, True], ids=["unenforced", "foreign_keys"])
+async def insert_memory(
+ request: pytest.FixtureRequest, sqlite_instance: SQLiteMemory, tmp_path: Path
+) -> AsyncIterator[MemoryInterface]:
+ if not request.param:
+ yield sqlite_instance
+ return
+ async with _foreign_key_memory_async(tmp_path / "foreign-keys.db") as memory:
+ yield memory
+
+
+@pytest.mark.parametrize("path", ["attack", "new_attack", "scenario", "scenario_update", "finding"])
+async def test_activity_is_rejected_when_its_operation_is_gone_at_insert_async(
+ insert_memory: MemoryInterface, path: str
+) -> None:
+ # Skipping the first check stands in for a concurrent delete between that check and the insert.
+ # Without foreign keys the post-insert check catches it; with them, as on SQL Server, the insert fails first.
+ operation_id = uuid4()
+ attack = AttackResult(conversation_id="orphan", objective="objective", operation_id=operation_id)
+ scenario = _scenario()
+ scenario.operation_id = operation_id
+ with patch.object(MemoryInterface, "_validate_activity_operation"), pytest.raises(LookupError, match="Operation"):
+ if path == "attack":
+ await insert_memory.add_attack_results_to_memory_async(attack_results=[attack])
+ elif path == "new_attack":
+ await insert_memory.add_conversation_branches_to_attack_async(
+ attack_result_id=attack.attack_result_id,
+ conversations=[Conversation(conversation_id=attack.conversation_id)],
+ message_pieces=[],
+ new_attack=attack,
+ )
+ elif path == "scenario":
+ await insert_memory.add_scenario_results_to_memory_async(scenario_results=[scenario])
+ elif path == "scenario_update":
+ await insert_memory.update_scenario_result_async(scenario_result=scenario)
+ else:
+ await insert_memory.add_finding_async(
+ finding=Finding(operation_id=operation_id, title="Orphan", severity="low")
+ )
+ assert not await insert_memory.get_attack_results_async()
+ assert not await insert_memory.get_scenario_results_async()
+ assert not await insert_memory.get_findings_async()
+
+
+async def test_unrelated_integrity_errors_still_propagate_with_foreign_keys_async(
+ foreign_key_memory: SQLiteMemory,
+) -> None:
+ operation = await foreign_key_memory.add_operation_async(Operation(name="Saved"))
+ attack = AttackResult(conversation_id="twice", objective="objective", operation_id=operation.id)
+ await foreign_key_memory.add_attack_results_to_memory_async(attack_results=[attack])
+ with pytest.raises(IntegrityError):
+ await foreign_key_memory.add_attack_results_to_memory_async(attack_results=[attack])
+
+
+async def test_delete_rejected_by_a_foreign_key_reports_the_usage_async(sqlite_instance: MemoryInterface) -> None:
+ # SQL Server's foreign keys fail the delete when a reference commits after the NOT EXISTS checks ran.
+ operation = await sqlite_instance.add_operation_async(Operation(name="Engagement"))
+ await _refer_to_operation_async(sqlite_instance, operation_id=operation.id, reference="attack")
+ execute = AsyncSession.execute
+
+ async def reject_delete_async(self: AsyncSession, statement: Any, *args: Any, **kwargs: Any) -> Any:
+ if isinstance(statement, sa.Delete):
+ raise IntegrityError("DELETE", {}, Exception("REFERENCE constraint"))
+ return await execute(self, statement, *args, **kwargs)
+
+ with patch.object(AsyncSession, "execute", reject_delete_async), pytest.raises(OperationInUseError) as raised:
+ await sqlite_instance.delete_operation_async(operation_id=operation.id)
+ assert raised.value.usage == OperationUsage(findings=0, attacks=1, scenario_runs=0)
diff --git a/tests/unit/memory/test_findings.py b/tests/unit/memory/test_findings.py
index 178646e994..9a2d095c28 100644
--- a/tests/unit/memory/test_findings.py
+++ b/tests/unit/memory/test_findings.py
@@ -1,11 +1,14 @@
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
+import sqlite3
from datetime import UTC, datetime
-from unittest.mock import MagicMock, patch
+from pathlib import Path
+from unittest.mock import AsyncMock, MagicMock, patch
from uuid import UUID, uuid4
import pytest
+from pydantic import ValidationError
from sqlalchemy import func, insert, select
from sqlalchemy.dialects import mssql
from sqlalchemy.engine import ScalarResult
@@ -17,6 +20,7 @@
AttackResultEntry,
FindingEntry,
FindingEvidenceEntry,
+ OperationEntry,
PromptMemoryEntry,
ScoreEntry,
)
@@ -35,13 +39,71 @@ async def test_operations_are_unique_by_trimmed_case_folded_name(sqlite_instance
assert await sqlite_instance.require_operation_async(operation_id=saved.id) == saved
-async def test_update_operation_replaces_only_the_description(sqlite_instance: SQLiteMemory) -> None:
+async def test_update_operation_renames_and_replaces_the_description(sqlite_instance: SQLiteMemory) -> None:
saved = await sqlite_instance.add_operation_async(Operation(name="Scoped", description="First draft"))
- updated = await sqlite_instance.update_operation_async(operation_id=saved.id, description="Final scope")
- assert updated == saved.model_copy(update={"description": "Final scope"})
- assert await sqlite_instance.require_operation_async(operation_id=saved.id) == updated
+ described = await sqlite_instance.update_operation_async(operation_id=saved.id, description="Final scope")
+ assert described == saved.model_copy(update={"description": "Final scope"})
+ recased = await sqlite_instance.update_operation_async(operation_id=saved.id, name=" SCOPED ")
+ assert recased == described.model_copy(update={"name": "SCOPED"})
+ renamed = await sqlite_instance.update_operation_async(operation_id=saved.id, name="Engagement", description="")
+ assert renamed == saved.model_copy(update={"name": "Engagement", "description": ""})
+ assert await sqlite_instance.get_operations_async(name="scoped") == []
+ assert await sqlite_instance.get_operations_async(name="engagement") == [renamed]
+
+
+async def test_update_operation_keeps_names_unique_and_valid(sqlite_instance: SQLiteMemory) -> None:
+ saved = await sqlite_instance.add_operation_async(Operation(name="Scoped"))
+ other = await sqlite_instance.add_operation_async(Operation(name="Other"))
+ taken = await sqlite_instance.update_operation_async(operation_id=saved.id, name=" OTHER ", description="Lost")
+ assert taken == other
+ assert await sqlite_instance.require_operation_async(operation_id=saved.id) == saved
+ with pytest.raises(ValidationError):
+ await sqlite_instance.update_operation_async(operation_id=saved.id, name=" ")
with pytest.raises(LookupError):
await sqlite_instance.update_operation_async(operation_id=uuid4(), description="Missing")
+ assert await sqlite_instance.get_operations_async() == [other, saved]
+
+
+async def test_add_finding_rejects_a_missing_operation(sqlite_instance: SQLiteMemory) -> None:
+ with pytest.raises(LookupError, match="Operation"):
+ await sqlite_instance.add_finding_async(finding=Finding(operation_id=uuid4(), title="Orphan", severity="low"))
+ assert await sqlite_instance.get_findings_async() == []
+
+
+async def test_add_finding_maps_a_foreign_key_failure_to_a_missing_operation(sqlite_instance: SQLiteMemory) -> None:
+ # SQL Server rejects the insert at flush when the Operation was deleted after the service checked it.
+ failure = AsyncMock(side_effect=IntegrityError("INSERT", {}, Exception("FOREIGN KEY constraint failed")))
+ saved = await sqlite_instance.add_operation_async(Operation(name="Saved"))
+ with patch.object(AsyncSession, "flush", failure), pytest.raises(LookupError, match="Operation"):
+ await sqlite_instance.add_finding_async(finding=Finding(operation_id=uuid4(), title="Orphan", severity="low"))
+ with patch.object(AsyncSession, "flush", failure), pytest.raises(IntegrityError):
+ await sqlite_instance.add_finding_async(finding=Finding(operation_id=saved.id, title="Other", severity="low"))
+
+
+async def test_update_operation_deleted_after_it_was_read_is_not_found(tmp_path: Path) -> None:
+ db_path = tmp_path / "update-race.db"
+ # A second connection needs a real file, and SQLiteMemory(...) would return the fixture's singleton.
+ memory = SQLiteMemory.__new__(SQLiteMemory)
+ with patch.object(memory, "cleanup"):
+ memory.__init__(db_path=db_path, silent=True)
+ try:
+ operation = await memory.add_operation_async(Operation(name="Racing"))
+ read = OperationEntry.get_operation
+
+ def read_then_delete(entry: OperationEntry) -> Operation:
+ other = sqlite3.connect(db_path, timeout=5)
+ try:
+ other.execute("DELETE FROM OperationEntries")
+ other.commit()
+ finally:
+ other.close()
+ return read(entry)
+
+ with patch.object(OperationEntry, "get_operation", read_then_delete), pytest.raises(LookupError):
+ await memory.update_operation_async(operation_id=operation.id, name="Renamed")
+ assert await memory.get_operations_async() == []
+ finally:
+ await memory.dispose_engine_async()
@pytest.mark.parametrize(