diff --git a/.pyrit_conf_example b/.pyrit_conf_example index 51488ee0d8..4fbad18e24 100644 --- a/.pyrit_conf_example +++ b/.pyrit_conf_example @@ -161,6 +161,21 @@ allow_custom_initializers: false # as scenario_presets_source. Defaults to ~/.pyrit/scenario_suites. # scenario_suites_source: https://account.blob.core.windows.net/container/scenario_suites +# Optional storage for saved targets, converters, scorers, and techniques, so they +# come back when the backend restarts. Each kind has its own folder under the source. +# This supports the same local directory and Azure Blob sources as +# scenario_presets_source. If any saved entry can't be restored, the backend stays +# not ready and lists the failures in its runtime status. Fix or delete the listed +# files, then apply the configuration or restart the backend. Defaults to +# ~/.pyrit/registry_recipes. +# +# ⚠️ WARNING: The backend rebuilds every saved entry in this source as an +# administrator, at startup and live reinitialization. Anyone who can write to this +# source can make the backend run a model's repository code or use its identity. +# Restrict write access to trusted administrators, and use a dedicated folder or +# blob prefix. +# registry_recipes_source: https://account.blob.core.windows.net/container/registry_recipes + # Upload Folder for API-Created Targets # ------------------------------------- # Folder that targets created through the REST API may upload local files from diff --git a/doc/getting_started/pyrit_conf.md b/doc/getting_started/pyrit_conf.md index bba64b45d4..64b8dc2b24 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -194,6 +194,20 @@ Suites store ordered references to scenario presets and resolve each preset's cu A SAS query string may be included; otherwise, PyRIT uses `DefaultAzureCredential`. The default is `~/.pyrit/scenario_suites`. Local writes use the same hidden lock-file and atomic replacement behavior as scenario presets. Startup scans the configured source so malformed suite documents are reported immediately without preventing the backend from starting. +### `registry_recipes_source` + +Stores saved targets, converters, scorers, and techniques in a local directory or Azure Blob container, so they come back when the backend restarts. An Azure URI may include a blob-name prefix: + +```yaml +registry_recipes_source: https://account.blob.core.windows.net/pyrit-storage/registry-recipes +``` + +Each kind has its own folder under the source (`targets/`, `converters/`, `scorers/`, `techniques/`), with one JSON file per saved entry named after the entry. A file holds the create request the entry was made with, and restore builds the entry through the API's create for its kind; restore refuses a file that stores a value for a parameter its class marks as a credential, such as `api_key`. A SAS query string may be included; otherwise, PyRIT uses `DefaultAzureCredential`. The default is `~/.pyrit/registry_recipes`. + +A saved file doesn't record who saved it, so the backend rebuilds every entry in this source as an administrator, at startup and after a live apply. A file here can, for example, make the backend run a model's repository code, or use the server's identity on a storage container that the file names. Write access to this source is administrator authority on the backend: grant it only to principals you already trust as backend administrators, and use a dedicated directory or blob prefix that less-trusted users and services cannot write to. + +The backend restores saved entries at startup and after a live apply, once the configured initializers have registered theirs. If any saved entry can't be restored, the backend doesn't become ready (state `failed`, outcome `restore-failed`), so a scan can't quietly run without an entry it expects. Until then, only the configuration and status routes answer; the rest return 503. Administrators see each failure with its reason in `GET /api/config/runtime` and on the GUI's Configuration page; everyone else sees that an administrator has to fix saved entries. Fix or delete the listed files, or, when a folder couldn't be listed, check that the backend can read this source. Then apply the configuration, if `enable_live_reinitialization` is on and the backend runs as one worker and one replica, which live apply requires; otherwise, restart the backend. Files whose names aren't valid entry names (up to 64 lowercase letters, digits, and underscores, starting with a letter) aren't failures and don't block startup; administrators can see them under `restore_ignored` in `GET /api/config/runtime`. + ### `initialization_scripts` Local paths to custom Python scripts containing `PyRITInitializer` subclasses. Paths can be absolute or relative to the current working directory. diff --git a/frontend/e2e/fixtures/manual_send_backend.py b/frontend/e2e/fixtures/manual_send_backend.py index 0c46ed2398..8be9b44941 100644 --- a/frontend/e2e/fixtures/manual_send_backend.py +++ b/frontend/e2e/fixtures/manual_send_backend.py @@ -4,6 +4,7 @@ """Offline backend for manual-send browser tests, with no real configuration or providers.""" import os +import tempfile from collections.abc import AsyncGenerator from contextlib import asynccontextmanager from unittest.mock import AsyncMock, patch @@ -27,19 +28,22 @@ async def _initialize_offline_async(self: ConfigurationLoader, *, raise_on_initi @asynccontextmanager async def _offline_lifespan_async(application: FastAPI) -> AsyncGenerator[None, None]: - config = ConfigurationLoader(memory_db_type="in_memory", env_files=[], env_akv_ref=[]) - with ( - patch.dict(os.environ, {"PYRIT_AUTH_MODE": "local", "PYRIT_DEV_MODE": "true"}, clear=True), - patch.object( - ConfigurationFileService, - "_read_source_async", - new=AsyncMock(return_value="memory_db_type: in_memory\nenv_files: []\ninitializers: []\n"), - ), - patch.object(ConfigurationLoader, "load_with_overrides", return_value=config), - patch.object(ConfigurationLoader, "initialize_pyrit_async", new=_initialize_offline_async), - ): - async with lifespan(application): - yield + with tempfile.TemporaryDirectory() as recipes: + config = ConfigurationLoader( + memory_db_type="in_memory", env_files=[], env_akv_ref=[], registry_recipes_source=recipes + ) + with ( + patch.dict(os.environ, {"PYRIT_AUTH_MODE": "local", "PYRIT_DEV_MODE": "true"}, clear=True), + patch.object( + ConfigurationFileService, + "_read_source_async", + new=AsyncMock(return_value="memory_db_type: in_memory\nenv_files: []\ninitializers: []\n"), + ), + patch.object(ConfigurationLoader, "load_with_overrides", return_value=config), + patch.object(ConfigurationLoader, "initialize_pyrit_async", new=_initialize_offline_async), + ): + async with lifespan(application): + yield app.router.lifespan_context = _offline_lifespan_async diff --git a/frontend/src/App.test.tsx b/frontend/src/App.test.tsx index 7519b77bd8..5c423233c3 100644 --- a/frontend/src/App.test.tsx +++ b/frontend/src/App.test.tsx @@ -4,6 +4,7 @@ */ import { Suspense, useEffect, useLayoutEffect, useRef, useState } from "react"; +import { fluentProviderClassNames } from "@fluentui/react-components"; import { act, render, screen, fireEvent, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { MemoryRouter, parsePath, useLocation, useNavigate } from "react-router"; @@ -139,6 +140,7 @@ jest.mock("./components/Layout/MainLayout", () => { onLabelsChange, operatorReadOnly, featureNavigation, + banner, }: { children: React.ReactNode; currentView: string; @@ -147,10 +149,12 @@ jest.mock("./components/Layout/MainLayout", () => { onLabelsChange: (labels: Record) => void; operatorReadOnly: boolean; featureNavigation?: React.ReactNode; + banner?: React.ReactNode; }) => { const navigate = useNavigate(); return (
+ {banner} {JSON.stringify(labels)} {JSON.stringify(labels)} {String(operatorReadOnly)} @@ -870,6 +874,18 @@ describe("App", () => { expect(screen.getByTestId("home-view")).toBeInTheDocument(); }); + it("says inside the app's layout and theme that saved entries couldn't be restored", async () => { + const { runtimeApi } = jest.requireMock("./services/api"); + runtimeApi.getReadiness.mockResolvedValueOnce({ + ready: false, state: "failed", generation: "first", restore_failures: 1, + }); + renderApp(); + + const message = await screen.findByText(/Some saved entries couldn't be restored/); + expect(screen.getByTestId("main-layout")).toContainElement(message); + expect(message.closest(`.${fluentProviderClassNames.root}`)).not.toBeNull(); + }); + it("renders the view named by the initial URL", () => { renderApp("/registry/targets"); diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx index 2ccb3e6bbe..bb0ff53a1e 100644 --- a/frontend/src/App.tsx +++ b/frontend/src/App.tsx @@ -869,6 +869,7 @@ function AppContent({ operatorAlias }: { operatorAlias: string | null }) { } currentView={currentView} onNavigate={handleNavigate} onOpenFeedback={() => setFeedbackOpen(true)} @@ -1106,7 +1107,6 @@ function App() { const operatorAlias = account?.username ? (account.username.split('@')[0]?.toLowerCase() ?? null) : null return ( - diff --git a/frontend/src/components/Configuration/Reinitialize.test.tsx b/frontend/src/components/Configuration/Reinitialize.test.tsx index 78614257fd..a40635f3fc 100644 --- a/frontend/src/components/Configuration/Reinitialize.test.tsx +++ b/frontend/src/components/Configuration/Reinitialize.test.tsx @@ -124,6 +124,86 @@ describe('Reinitialize', () => { expect(screen.getByRole('button', { name: 'Reinitialize PyRIT' })).toBeDisabled() }) + it('should let a failed restore be retried by reinitializing', async () => { + const user = userEvent.setup() + api.getRuntimeStatus.mockResolvedValue({ + ...ready, + state: 'failed', + outcome: 'restore-failed', + message: "1 saved registry entry couldn't be restored.", + }) + render( + + + , + ) + expect(await screen.findByText(/Runtime: failed\. 1 saved registry entry/)).toBeInTheDocument() + await waitFor(() => expect(screen.getByRole('button', { name: 'Reinitialize PyRIT' })).toBeEnabled()) + await user.click(screen.getByRole('button', { name: 'Reinitialize PyRIT' })) + await user.click(within(await screen.findByRole('dialog')).getByRole('button', { name: 'Reinitialize PyRIT' })) + await waitFor(() => expect(api.reinitialize).toHaveBeenCalledWith('saved-v1')) + }) + + it('should list the saved entries that could not be restored', async () => { + api.getRuntimeStatus.mockResolvedValue({ + ...ready, + state: 'failed', + outcome: 'restore-failed', + message: "3 saved registry entries couldn't be restored.", + restore_failures: [ + { + kind: 'targets', + name: 'pool', + reason: 'missing_dependency', + message: "Uses target 'gone', which doesn't exist.", + version: 'v1', + dependency: { kind: 'targets', name: 'gone' }, + }, + { + kind: 'scorers', + name: 'refusal', + reason: 'failed_dependency', + message: "It uses an entry that couldn't be restored.", + version: 'v2', + dependency: { kind: 'targets', name: 'pool' }, + }, + { + kind: 'converters', + name: null, + reason: 'unreadable_storage', + message: "The saved converters couldn't be listed: permission denied", + version: null, + dependency: null, + }, + ], + }) + render( + + + , + ) + + const region = await screen.findByRole('region', { name: "Saved entries that couldn't be restored" }) + expect(within(region).getByText("Target pool: Uses target 'gone', which doesn't exist.")).toBeInTheDocument() + expect(within(region).getByText("Scorer refusal: It uses an entry that couldn't be restored. (target pool)")) + .toBeInTheDocument() + expect(within(region).getByText("Saved converters: The saved converters couldn't be listed: permission denied")) + .toBeInTheDocument() + expect(region).toHaveAttribute('tabindex', '0') + }) + + it('should show no list when every saved entry was restored', async () => { + api.getRuntimeStatus.mockResolvedValue({ ...ready, restore_failures: [] }) + render( + + + , + ) + + await waitFor(() => expect(screen.getByRole('button', { name: 'Reinitialize PyRIT' })).toBeEnabled()) + expect(screen.queryByRole('region', { name: "Saved entries that couldn't be restored" })).not.toBeInTheDocument() + }) + describe('runtime status polling', () => { beforeEach(() => { jest.useFakeTimers() diff --git a/frontend/src/components/Configuration/Reinitialize.tsx b/frontend/src/components/Configuration/Reinitialize.tsx index 193a661459..2d17fc0d50 100644 --- a/frontend/src/components/Configuration/Reinitialize.tsx +++ b/frontend/src/components/Configuration/Reinitialize.tsx @@ -13,6 +13,7 @@ import { toApiError } from '@/services/errors' import type { RuntimeStatus } from '@/types' import { useReinitializeStyles } from './Reinitialize.styles' +import RestoreFailures from './RestoreFailures' interface ReinitializeProps { version: string @@ -81,13 +82,17 @@ export default function Reinitialize({ return (
{status && (showRuntimeStatus || statusIsError || !status.enabled) && ( - + // Multiline, since a status message such as a restore failure can be long and change after the bar lays out. + Runtime: {status.state}. {status.message} {!status.enabled && ' Reinitialization requires one backend worker and one replica.'} )} + {status?.restore_failures && status.restore_failures.length > 0 && ( + + )} {pollError && {pollError}} {applyError && {applyError}} {hasUnsavedChanges && Save or explicitly discard unsaved edits before reinitializing.} diff --git a/frontend/src/components/Configuration/RestoreFailures.styles.ts b/frontend/src/components/Configuration/RestoreFailures.styles.ts new file mode 100644 index 0000000000..cacf3e052f --- /dev/null +++ b/frontend/src/components/Configuration/RestoreFailures.styles.ts @@ -0,0 +1,15 @@ +import { makeStyles } from '@fluentui/react-components' + +export const useRestoreFailuresStyles = makeStyles({ + // The page doesn't scroll past the editor, so a long list scrolls on its own. It fills its column, and a long + // name or message wraps instead of widening the list past it. + list: { + alignSelf: 'stretch', + boxSizing: 'border-box', + maxWidth: '100%', + minWidth: 0, + maxHeight: '30vh', + overflowWrap: 'anywhere', + overflowY: 'auto', + }, +}) diff --git a/frontend/src/components/Configuration/RestoreFailures.tsx b/frontend/src/components/Configuration/RestoreFailures.tsx new file mode 100644 index 0000000000..7cf08c6b4b --- /dev/null +++ b/frontend/src/components/Configuration/RestoreFailures.tsx @@ -0,0 +1,42 @@ +import type { RegistryRecipeKind, RestoreFailure } from '@/types' + +import { useRestoreFailuresStyles } from './RestoreFailures.styles' + +const KIND_LABELS: Record = { + targets: 'Target', + converters: 'Converter', + scorers: 'Scorer', + techniques: 'Technique', +} + +function describeFailure(failure: RestoreFailure): string { + const subject = failure.name === null ? `Saved ${failure.kind}` : `${KIND_LABELS[failure.kind]} ${failure.name}` + // The message usually names the entry it depends on already. + const dependency = failure.dependency && !failure.message.includes(`'${failure.dependency.name}'`) + ? ` (${KIND_LABELS[failure.dependency.kind].toLowerCase()} ${failure.dependency.name})` + : '' + return `${subject}: ${failure.message}${dependency}` +} + +interface RestoreFailuresProps { + failures: RestoreFailure[] +} + +export default function RestoreFailures({ failures }: RestoreFailuresProps) { + const styles = useRestoreFailuresStyles() + return ( +
+
    + {failures.map((failure) => ( +
  • {describeFailure(failure)}
    • + ))} +
+
+ ) +} diff --git a/frontend/src/components/Layout/MainLayout.styles.ts b/frontend/src/components/Layout/MainLayout.styles.ts index ab38ec00d2..2e740fbfc2 100644 --- a/frontend/src/components/Layout/MainLayout.styles.ts +++ b/frontend/src/components/Layout/MainLayout.styles.ts @@ -33,6 +33,9 @@ export const useMainLayoutStyles = makeStyles({ transform: 'translateY(0)', }, }, + banner: { + flexShrink: 0, + }, topBar: { minHeight: '60px', flexShrink: 0, diff --git a/frontend/src/components/Layout/MainLayout.test.tsx b/frontend/src/components/Layout/MainLayout.test.tsx index 1fa3da9fc7..4f36c2f90a 100644 --- a/frontend/src/components/Layout/MainLayout.test.tsx +++ b/frontend/src/components/Layout/MainLayout.test.tsx @@ -77,6 +77,23 @@ describe("MainLayout", () => { expect(screen.getByRole("button", { name: "Collapse sidebar" })).toHaveAttribute("aria-expanded", "true"); }); + it("shows a banner between the skip link and the top bar, in the layout column", async () => { + mockedVersionApi.getVersion.mockResolvedValue({ version: "1.0.0" }); + + renderWithProvider( + Saved entries need attention
}> +
Child Content
+ + ); + + const skipLink = screen.getByRole("link", { name: "Skip to main content" }); + const banner = screen.getByRole("status"); + const subtitle = screen.getByText("Python Risk Identification Tool"); + expect(skipLink.compareDocumentPosition(banner) & Node.DOCUMENT_POSITION_FOLLOWING).toBeTruthy(); + expect(banner.compareDocumentPosition(subtitle) & Node.DOCUMENT_POSITION_FOLLOWING).toBeTruthy(); + expect(screen.getByRole("main")).not.toContainElement(banner); + }); + it("renders the header with title and subtitle", async () => { mockedVersionApi.getVersion.mockResolvedValue({ version: "1.0.0" }); diff --git a/frontend/src/components/Layout/MainLayout.tsx b/frontend/src/components/Layout/MainLayout.tsx index 7ae6bda593..0ef35898f0 100644 --- a/frontend/src/components/Layout/MainLayout.tsx +++ b/frontend/src/components/Layout/MainLayout.tsx @@ -35,6 +35,8 @@ interface MainLayoutProps extends OperationControlProps { onStartTour?: () => void labelsBusy?: boolean featureNavigation?: React.ReactNode + /** Shown above the top bar, so the content below gets smaller instead of losing its bottom. */ + banner?: React.ReactNode } export default function MainLayout({ @@ -53,6 +55,7 @@ export default function MainLayout({ onStartTour, labelsBusy = false, featureNavigation, + banner, }: MainLayoutProps) { usePreventFileNavigation() const styles = useMainLayoutStyles() @@ -89,6 +92,7 @@ export default function MainLayout({ Skip to main content + {banner &&
{banner}
}
({ runtimeApi: { getReadiness: jest.fn() } })) +jest.mock('@/services/api', () => ({ + runtimeApi: { getReadiness: jest.fn() }, +})) + +function TestWrapper({ children }: { children: ReactNode }) { + return {children} +} function Consumer() { const runtime = useRuntime() @@ -41,3 +49,59 @@ describe('RuntimeProvider', () => { expect(screen.getByRole('textbox', { name: 'Draft' })).toHaveValue('keep my draft') }) }) + +function renderBanner() { + render() +} + +describe('RuntimeBanner', () => { + beforeEach(() => { jest.clearAllMocks() }) + + it('should tell everyone that an administrator has to fix saved entries that could not be restored', async () => { + jest.mocked(runtimeApi.getReadiness).mockResolvedValue({ + ready: false, + state: 'failed', + generation: 'first', + restore_failures: 2, + }) + + renderBanner() + + expect(await screen.findByText(new RegExp( + "Some saved entries couldn't be restored, so new runs and sends are unavailable\\. An administrator has to " + + 'view the failures in Configuration, fix or delete the files in registry_recipes_source, then reinitialize ' + + 'PyRIT or restart the backend\\. Drafts are preserved\\.', + ))).toBeInTheDocument() + expect(screen.queryByRole('region')).not.toBeInTheDocument() + expect(screen.queryByRole('button')).not.toBeInTheDocument() + expect(screen.queryByText(/PyRIT runtime: failed/)).not.toBeInTheDocument() + }) + + it('should show the general runtime bar while the runtime is not ready for another reason', async () => { + jest.mocked(runtimeApi.getReadiness).mockResolvedValue({ + ready: false, + state: 'restart-required', + generation: '', + restore_failures: 1, + }) + + renderBanner() + + expect(await screen.findByText(/PyRIT runtime: restart-required/)).toBeInTheDocument() + expect(screen.queryByText(/couldn't be restored/)).not.toBeInTheDocument() + }) + + it('should show nothing while the runtime is ready', async () => { + jest.mocked(runtimeApi.getReadiness).mockResolvedValue({ + ready: true, + state: 'ready', + generation: 'first', + restore_failures: 0, + }) + + renderBanner() + + await waitFor(() => expect(screen.queryByText(/PyRIT runtime/)).not.toBeInTheDocument()) + expect(screen.queryByText(/couldn't be restored/)).not.toBeInTheDocument() + }) +}) diff --git a/frontend/src/hooks/useRuntime.tsx b/frontend/src/hooks/useRuntime.tsx index 2035358d63..485ca25e03 100644 --- a/frontend/src/hooks/useRuntime.tsx +++ b/frontend/src/hooks/useRuntime.tsx @@ -41,8 +41,21 @@ export function RuntimeProvider({ children }: { children: ReactNode }) { export function RuntimeBanner() { const runtime = useRuntime() if (runtime.ready) return null + // Multiline, since the text changes after the bar first lays out, and Fluent's automatic layout only rechecks + // on a resize. + if (runtime.state === 'failed' && (runtime.restore_failures ?? 0) > 0) { + return ( + + + Some saved entries couldn't be restored, so new runs and sends are unavailable. An administrator has to + view the failures in Configuration, fix or delete the files in registry_recipes_source, then reinitialize + PyRIT or restart the backend. Drafts are preserved. + + + ) + } return ( - + PyRIT runtime: {runtime.state}. New runs and sends are unavailable. Administrators can repair saved sources and retry from Configuration. Drafts are preserved. diff --git a/frontend/src/types/index.ts b/frontend/src/types/index.ts index 3f3331a30f..46b0e5aed4 100644 --- a/frontend/src/types/index.ts +++ b/frontend/src/types/index.ts @@ -2496,6 +2496,25 @@ export interface RuntimeReadiness { ready: boolean state: string generation: string + /** How many saved entries couldn't be restored; administrators see why in the runtime status. */ + restore_failures?: number +} + +export type RegistryRecipeKind = 'targets' | 'converters' | 'scorers' | 'techniques' + +export interface SavedEntryRef { + kind: RegistryRecipeKind + name: string +} + +/** A saved entry, or a whole folder of them when `name` is null, that couldn't be restored. */ +export interface RestoreFailure { + kind: RegistryRecipeKind + name: string | null + reason: string + message: string + version: string | null + dependency: SavedEntryRef | null } export interface RuntimeStatus { @@ -2506,6 +2525,8 @@ export interface RuntimeStatus { applying: boolean outcome: string message: string + /** The saved entries that couldn't be restored; only administrators can read the runtime status. */ + restore_failures?: RestoreFailure[] } export type FindingSeverity = 'critical' | 'important' | 'moderate' | 'low' | 'informational' | 'other' diff --git a/pyrit/backend/README.md b/pyrit/backend/README.md index 444cd2be87..ab3aae9ee6 100644 --- a/pyrit/backend/README.md +++ b/pyrit/backend/README.md @@ -691,8 +691,9 @@ The chat preserves failed drafts across replacement, but requires converter choi to be reviewed again rather than restoring outputs tied to the previous registry. The offline browser fixture `frontend.e2e.fixtures.manual_send_backend:app` uses the real -backend lifecycle with isolated in-memory SQLite, no environment files, and no default -providers. The recovery tests register deterministic loopback targets. For example, start +backend lifecycle with isolated in-memory SQLite, no environment files, no saved registry +entries, and no default providers. The recovery tests register deterministic loopback +targets. For example, start it with `uv run python -m uvicorn frontend.e2e.fixtures.manual_send_backend:app --host 127.0.0.1 --port 18213` (set `PYTHONUTF8=1` on Windows). In another shell, set `PYRIT_BACKEND_URL=http://127.0.0.1:18213` and `E2E_FRONTEND_PORT=31213`, then run `npx playwright test chat-recovery --project seeded --workers 1` diff --git a/pyrit/backend/middleware/error_handlers.py b/pyrit/backend/middleware/error_handlers.py index 0beaff62cf..10d7c2c481 100644 --- a/pyrit/backend/middleware/error_handlers.py +++ b/pyrit/backend/middleware/error_handlers.py @@ -6,6 +6,7 @@ """ import logging +import traceback from fastapi import FastAPI, Request, status from fastapi.exceptions import RequestValidationError @@ -13,8 +14,10 @@ from pyrit.backend.middleware.request_size import RequestTooLargeError, request_too_large_response from pyrit.backend.models import FieldError, ProblemDetail +from pyrit.common.azure_storage import redact_urls from pyrit.exceptions.analytics_exception import AnalyticsException from pyrit.exceptions.dataset_exception import DatasetAPIException +from pyrit.registry.file_document_storage import DocumentStorageUnavailableError, DocumentUnreadableError logger = logging.getLogger(__name__) @@ -192,6 +195,41 @@ async def permission_error_handler( # pyrit-async-suffix-exempt content=problem.model_dump(exclude_none=True), ) + @app.exception_handler(DocumentStorageUnavailableError) + async def storage_unavailable_handler( # pyrit-async-suffix-exempt + request: Request, + exc: DocumentStorageUnavailableError, + ) -> JSONResponse: # pyrit-keyword-only-exempt Starlette exception handler (request, exc) + """ + Handle stored documents that can't be read or written as 503 Service Unavailable. + + The log names the storage location, without credentials, and the reason, with the traceback at debug + level. URL credentials are removed from the traceback, since the error behind it can carry a SAS URL. + The response never carries a path or URL. + + Returns: + JSONResponse: RFC 7807 problem detail response with 503 status. + """ + logger.warning(f"{exc} ({request.method} {request.url.path})") + if logger.isEnabledFor(logging.DEBUG): + logger.debug(f"Storage error details:\n{redact_urls(''.join(traceback.format_exception(exc)).rstrip())}") + if isinstance(exc, DocumentUnreadableError): + detail = f"{exc.label} '{exc.name}' can't be read: {exc.reason}" + else: + detail = f"{exc.label} storage is unavailable" + problem = ProblemDetail( + type="/errors/storage-unavailable", + title="Service Unavailable", + status=status.HTTP_503_SERVICE_UNAVAILABLE, + detail=detail, + instance=str(request.url.path), + ) + + return JSONResponse( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + content=problem.model_dump(exclude_none=True), + ) + @app.exception_handler(NotImplementedError) async def not_implemented_handler( # pyrit-async-suffix-exempt request: Request, diff --git a/pyrit/backend/routes/health.py b/pyrit/backend/routes/health.py index b6e06c91a4..e0cd361faa 100644 --- a/pyrit/backend/routes/health.py +++ b/pyrit/backend/routes/health.py @@ -13,30 +13,31 @@ router = APIRouter() -def _runtime_readiness(request: Request) -> dict[str, str | bool]: +def _runtime_readiness(request: Request) -> dict[str, str | bool | int]: runtime = getattr(request.app.state, "runtime_lifecycle", None) if runtime is None: - return {"ready": False, "state": "failed", "generation": ""} + return {"ready": False, "state": "failed", "generation": "", "restore_failures": 0} return { "ready": runtime.reported_state == "ready", "state": runtime.reported_state, "generation": runtime.generation, + "restore_failures": len(runtime.restore_failures), } @router.get("/runtime") -async def runtime_readiness_async(request: Request) -> dict[str, str | bool]: +async def runtime_readiness_async(request: Request) -> dict[str, str | bool | int]: """ - Expose readiness and generation without administrative details. + Expose readiness, generation, and how many saved entries couldn't be restored, without details. Returns: - dict[str, str | bool]: Lightweight runtime readiness. + dict[str, str | bool | int]: Lightweight runtime readiness. """ return _runtime_readiness(request) @router.get("/health") -async def health_check_async(*, request: Request, response: Response) -> dict[str, str | bool | None]: +async def health_check_async(*, request: Request, response: Response) -> dict[str, str | bool | int | None]: """ Check the health status of the backend service. @@ -48,7 +49,7 @@ async def health_check_async(*, request: Request, response: Response) -> dict[st deployment checks use the readiness fields instead. Returns: - dict[str, str | bool | None]: Health status, timestamp, runtime readiness, and serving ACA revision. + dict[str, str | bool | int | None]: Health status, timestamp, runtime readiness, and serving ACA revision. """ readiness = _runtime_readiness(request) response.headers["Cache-Control"] = "no-store" diff --git a/pyrit/backend/routes/initializers.py b/pyrit/backend/routes/initializers.py index 50fa79db85..529f92d641 100644 --- a/pyrit/backend/routes/initializers.py +++ b/pyrit/backend/routes/initializers.py @@ -14,7 +14,6 @@ DELETE /api/initializers/{name} — unregister an initializer """ -from azure.core.exceptions import AzureError from fastapi import APIRouter, Depends, HTTPException, Query, Request, status from pyrit.backend.middleware.auth import require_admin @@ -35,19 +34,6 @@ router = APIRouter(prefix="/initializers", tags=["initializers"]) -def _custom_storage_unavailable() -> HTTPException: - """ - Create a sanitized response for unavailable custom initializer storage. - - Returns: - HTTPException: A service-unavailable response without SDK details. - """ - return HTTPException( - status_code=status.HTTP_503_SERVICE_UNAVAILABLE, - detail="Custom initializer storage is temporarily unavailable", - ) - - def _configured_initializers(request: Request) -> list[ConfiguredInitializerSetting]: """ Read the initializer list captured from ``.pyrit_conf`` at backend startup. @@ -138,10 +124,7 @@ async def list_custom_initializers(request: Request) -> CustomInitializerListRes CustomInitializerListResponse: The configured source and stored scripts. """ _check_custom_initializers_allowed(request) - try: - return await get_initializer_service().list_custom_initializers_async() - except AzureError as exc: - raise _custom_storage_unavailable() from exc + return await get_initializer_service().list_custom_initializers_async() @router.get( @@ -205,8 +188,6 @@ async def register_initializer( # pyrit-async-suffix-exempt try: return await service.register_initializer_async(name=body.name, script_content=body.script_content) - except AzureError as exc: - raise _custom_storage_unavailable() from exc except ValueError as e: detail = str(e) if "already registered" in detail: @@ -243,8 +224,6 @@ async def unregister_initializer( # pyrit-async-suffix-exempt try: await service.unregister_initializer_async(initializer_name=initializer_name) - except AzureError as exc: - raise _custom_storage_unavailable() from exc except ValueError as e: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, diff --git a/pyrit/backend/routes/scenario_presets.py b/pyrit/backend/routes/scenario_presets.py index ad2880dd31..6d727491eb 100644 --- a/pyrit/backend/routes/scenario_presets.py +++ b/pyrit/backend/routes/scenario_presets.py @@ -19,7 +19,6 @@ from typing import Annotated -from azure.core.exceptions import AzureError from fastapi import APIRouter, Depends, HTTPException, Query, status from pyrit.backend.middleware.auth import current_user_name, require_admin @@ -38,19 +37,6 @@ router = APIRouter(prefix="/scenario-presets", tags=["scenario-presets"]) -def _storage_unavailable() -> HTTPException: - """ - Create a sanitized response for unavailable preset storage. - - Returns: - HTTPException: A service-unavailable response without SDK details. - """ - return HTTPException( - status_code=status.HTTP_503_SERVICE_UNAVAILABLE, - detail="Scenario preset storage is temporarily unavailable", - ) - - def _not_found(name: str) -> HTTPException: """ Create a not-found response for a missing preset. @@ -79,8 +65,6 @@ async def _load_preset_or_404_async(name: str) -> ScenarioPresetResponse: """ try: preset = await get_scenario_preset_service().get_preset_async(name=name) - except AzureError as exc: - raise _storage_unavailable() from exc except ValueError as exc: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None @@ -106,10 +90,7 @@ async def list_scenario_presets( # pyrit-async-suffix-exempt Returns: ScenarioPresetListResponse: The configured source and the stored presets. """ - try: - return await get_scenario_preset_service().list_presets_async(include_estimates=include_estimates) - except AzureError as exc: - raise _storage_unavailable() from exc + return await get_scenario_preset_service().list_presets_async(include_estimates=include_estimates) @router.get( @@ -177,8 +158,6 @@ async def create_scenario_preset( # pyrit-async-suffix-exempt status_code=status.HTTP_409_CONFLICT, detail=f"Scenario preset '{preset.name}' already exists", ) from exc - except AzureError as exc: - raise _storage_unavailable() from exc @router.put( @@ -239,8 +218,6 @@ async def update_scenario_preset( # pyrit-async-suffix-exempt status_code=status.HTTP_409_CONFLICT, detail=f"Scenario preset '{name}' changed since it was read; re-read it and retry", ) from exc - except AzureError as exc: - raise _storage_unavailable() from exc @router.delete( @@ -278,8 +255,6 @@ async def delete_scenario_preset( # pyrit-async-suffix-exempt status_code=status.HTTP_409_CONFLICT, detail=f"Scenario preset '{name}' changed since it was read; re-read it and retry", ) from None - except (AzureError, TimeoutError) as exc: - raise _storage_unavailable() from exc except ValueError as exc: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None @@ -329,8 +304,6 @@ async def resolve_scenario_preset( # pyrit-async-suffix-exempt status_code=status.HTTP_409_CONFLICT, detail=f"Scenario preset '{name}' changed since it was read; re-read it and retry", ) from None - except AzureError as exc: - raise _storage_unavailable() from exc except ValueError as exc: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None diff --git a/pyrit/backend/routes/scenario_suites.py b/pyrit/backend/routes/scenario_suites.py index 42501471c4..79ea691149 100644 --- a/pyrit/backend/routes/scenario_suites.py +++ b/pyrit/backend/routes/scenario_suites.py @@ -6,7 +6,6 @@ from typing import Annotated from uuid import UUID -from azure.core.exceptions import AzureError from fastapi import APIRouter, Depends, HTTPException, Query, Request, status from pyrit.backend.middleware.auth import current_user_name, require_admin @@ -37,19 +36,6 @@ router = APIRouter(prefix="/scenario-suites", tags=["scenario-suites"]) -def _storage_unavailable() -> HTTPException: - """ - Create a sanitized response for unavailable suite storage. - - Returns: - HTTPException: A service-unavailable response without SDK details. - """ - return HTTPException( - status_code=status.HTTP_503_SERVICE_UNAVAILABLE, - detail="Scenario suite storage is temporarily unavailable", - ) - - def _not_found(name: str) -> HTTPException: """ Create a not-found response for a missing suite. @@ -72,8 +58,6 @@ async def _load_suite_or_404_async(name: str) -> ScenarioSuiteResponse: """ try: suite = await get_scenario_suite_service().get_suite_async(name=name) - except AzureError as exc: - raise _storage_unavailable() from exc except ValueError as exc: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None @@ -90,10 +74,7 @@ async def list_scenario_suites() -> ScenarioSuiteListResponse: # pyrit-async-su Returns: ScenarioSuiteListResponse: Suites from both origins with current entry statuses. """ - try: - return await get_scenario_suite_service().list_suites_async() - except AzureError as exc: - raise _storage_unavailable() from exc + return await get_scenario_suite_service().list_suites_async() @router.post("/{name}/readiness", response_model=ScenarioReadiness) @@ -116,8 +97,6 @@ async def get_scenario_suite_readiness_async( raise HTTPException(status_code=409, detail=str(exc)) from None except ValueError as exc: raise HTTPException(status_code=400, detail=str(exc)) from None - except AzureError as exc: - raise _storage_unavailable() from exc @router.get( @@ -211,8 +190,6 @@ async def launch_scenario_suite( # pyrit-async-suffix-exempt raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None except ScenarioSuiteExecutionCapacityError as exc: raise HTTPException(status_code=status.HTTP_429_TOO_MANY_REQUESTS, detail=str(exc)) from None - except AzureError as exc: - raise _storage_unavailable() from exc except RuntimeError as exc: raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail=str(exc)) from None @@ -269,8 +246,6 @@ async def create_scenario_suite( # pyrit-async-suffix-exempt status_code=status.HTTP_409_CONFLICT, detail=f"Scenario suite '{suite.name}' already exists", ) from exc - except AzureError as exc: - raise _storage_unavailable() from exc @router.put( @@ -319,8 +294,6 @@ async def update_scenario_suite( # pyrit-async-suffix-exempt status_code=status.HTTP_409_CONFLICT, detail=f"Scenario suite '{name}' changed since it was read; re-read it and retry", ) from exc - except AzureError as exc: - raise _storage_unavailable() from exc @router.delete( @@ -350,8 +323,6 @@ async def delete_scenario_suite( # pyrit-async-suffix-exempt status_code=status.HTTP_409_CONFLICT, detail=f"Scenario suite '{name}' changed since it was read; re-read it and retry", ) from None - except (AzureError, TimeoutError) as exc: - raise _storage_unavailable() from exc except ValueError as exc: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None diff --git a/pyrit/backend/services/runtime_lifecycle.py b/pyrit/backend/services/runtime_lifecycle.py index daef21f0ed..ffa1c7480e 100644 --- a/pyrit/backend/services/runtime_lifecycle.py +++ b/pyrit/backend/services/runtime_lifecycle.py @@ -7,7 +7,8 @@ import logging import os import uuid -from typing import Any +from functools import partial +from typing import TYPE_CHECKING, Any from fastapi import FastAPI @@ -15,6 +16,7 @@ from pyrit.backend.services.analytics_service import get_analytics_service from pyrit.backend.services.configuration_file_service import ConfigurationFileService from pyrit.backend.services.environment_file_service import EnvironmentFileService +from pyrit.backend.services.saved_entry_restore import build_saved_entry_async from pyrit.backend.services.scenario_run_service import get_scenario_run_service, peek_scenario_run_service from pyrit.backend.services.scenario_suite_execution_service import peek_scenario_suite_execution_service from pyrit.backend.services.service_lifecycle import ( @@ -26,10 +28,14 @@ from pyrit.common.path import CONFIGURATION_DIRECTORY_PATH from pyrit.memory import CentralMemory from pyrit.registry import InitializerRegistry, ScenarioPresetRegistry, ScenarioSuiteRegistry +from pyrit.registry.registry_recipe_restore import RegistryRecipeRestorer from pyrit.setup import ConfigurationLoader from pyrit.setup.environment_loading import resolve_environment_async from pyrit.setup.initialization import validate_reinitialization_memory +if TYPE_CHECKING: + from pyrit.models.catalog import IgnoredRecipeFile, RestoreFailure + logger = logging.getLogger(__name__) @@ -46,6 +52,8 @@ def __init__(self, *, app: FastAPI, source: ConfigurationFileService) -> None: self.version: str | None = None self.outcome = "starting" self.message = "" + self.restore_failures: list[RestoreFailure] = [] + self.restore_ignored: list[IgnoredRecipeFile] = [] self.operations: set[asyncio.Task[None]] = set() self.management_operations: set[asyncio.Task[None]] = set() self.apply_task: asyncio.Task[None] | None = None @@ -75,6 +83,8 @@ def status(self) -> dict[str, Any]: "message": self.message, "enabled": self.topology_supported, "applying": self.apply_task is not None and not self.apply_task.done(), + "restore_failures": [failure.model_dump(mode="json") for failure in self.restore_failures], + "restore_ignored": [file.model_dump(mode="json") for file in self.restore_ignored], } async def _load_async(self) -> ConfigurationLoader: @@ -119,6 +129,52 @@ def _publish(self, config: ConfigurationLoader) -> None: self.generation = str(uuid.uuid4()) self.state, self.outcome, self.message = "ready", "success", "PyRIT is ready." + async def _restore_saved_entries_async(self, config: ConfigurationLoader) -> bool: + """ + Rebuild saved registry entries after the startup scripts have registered theirs. + + A failure keeps the runtime not ready, since a scan that would have used a missing entry + could quietly fall back to a different one. Fixing or deleting the stored file, then applying + or restarting, recovers. + + Args: + config (ConfigurationLoader): The loaded configuration naming the storage source. + + Returns: + bool: Whether every saved entry was restored. + """ + restorer = RegistryRecipeRestorer(source=config.registry_recipes_source) + # Reading can wait on storage, so it runs in a worker thread. Entries are built on this loop, through + # the API's create for each kind, because a constructor such as HuggingFaceChatTarget's starts a task + # on it; that create does the target's slow torch setup in a worker thread first. + saved = await asyncio.to_thread(restorer.read) + report = await restorer.restore_async( + saved=saved, build_async=partial(build_saved_entry_async, upload_directory=config.target_upload_directory) + ) + self.restore_failures, self.restore_ignored = report.failures, report.ignored + if not report.failures: + return True + self.state, self.outcome = "failed", "restore-failed" + next_step = ( + "apply the configuration or restart the backend" + if config.enable_live_reinitialization and self.topology_supported + else "restart the backend" + ) + unread = sorted({failure.kind.value for failure in report.failures if failure.name is None}) + if unread: + self.message = ( + f"The saved {' and '.join(unread)} couldn't be read, so no saved entries were restored. " + f"Check that the backend can read registry_recipes_source, then {next_step}." + ) + else: + count = len(report.failures) + self.message = ( + f"{count} saved registry {'entry' if count == 1 else 'entries'} couldn't be restored. " + f"Fix or delete the files named in the server log or GET /api/config/runtime, then {next_step}." + ) + logger.error("PyRIT isn't ready: %s", self.message) + return False + async def startup_async(self) -> None: """Keep management alive even when parsing or initialization fails.""" self.app.state.allow_custom_initializers = False @@ -137,7 +193,8 @@ async def startup_async(self) -> None: await config.initialize_pyrit_async(raise_on_initializer_error=True) await get_scenario_run_service().reconcile_interrupted_runs_async() _, self.version = await self.source.read_with_version_async() - self._publish(config) + if await self._restore_saved_entries_async(config): + self._publish(config) except Exception: self.state, self.outcome = "restart-required", "initialization-failed" self.message = "PyRIT setup failed. Repair saved configuration, then restart the backend." @@ -158,7 +215,7 @@ def begin_apply(self, *, version: str) -> dict[str, Any]: "outcome": "unsupported", "message": ("Reinitialization requires one backend worker and one replica."), } - if self.state in ("failed", "restart-required"): + if self.state == "restart-required": return {**self.status(), "outcome": "restart-required"} if (self.apply_task and not self.apply_task.done()) or self.edit_lock.locked() or self.management_operations: return {**self.status(), "outcome": "busy"} @@ -179,7 +236,9 @@ async def _apply_async(self, *, version: str) -> None: if not config.enable_live_reinitialization: self.outcome = "unsupported" self.message = ( - "Set enable_live_reinitialization: true in the saved configuration to use live apply." + "Live apply is off. Restart the backend to restore the saved entries again." + if previous_state == "failed" + else "Set enable_live_reinitialization: true in the saved configuration to use live apply." ) return # Type compatibility is checked before contacting environment sources. @@ -215,11 +274,14 @@ async def _apply_async(self, *, version: str) -> None: return mutated = True + # The restore report describes the runtime being replaced. + self.restore_failures, self.restore_ignored = [], [] await close_services_async() await config.apply_prepared_reinitialization_async(prepared=prepared) await self._management_async(config) self.version = current_version - self._publish(config) + if await self._restore_saved_entries_async(config): + self._publish(config) except (Exception, asyncio.CancelledError, BaseExceptionGroup) as error: logger.exception( "PyRIT live apply failed phase=%s generation=%s", diff --git a/pyrit/backend/services/saved_entry_restore.py b/pyrit/backend/services/saved_entry_restore.py new file mode 100644 index 0000000000..a4ac353988 --- /dev/null +++ b/pyrit/backend/services/saved_entry_restore.py @@ -0,0 +1,86 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Build saved registry entries through the API's own create for each kind.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from pydantic import ValidationError + +from pyrit.backend.models import CreateConverterRequest, CreateScorerRequest, CreateTargetRequest +from pyrit.backend.models.techniques import CreateTechniqueRequest +from pyrit.backend.services.converter_service import get_converter_service +from pyrit.backend.services.scorer_service import get_scorer_service +from pyrit.backend.services.target_service import get_target_service +from pyrit.backend.services.technique_service import get_technique_service +from pyrit.models.catalog import RegistryRecipeKind, RestoreFailureReason, TechniqueRecipe +from pyrit.registry.registry_recipe_restore import RecipeRejectedError + +if TYPE_CHECKING: + from pyrit.models.catalog import StoredRecipe + + +async def build_saved_entry_async(stored: StoredRecipe, *, upload_directory: str | None) -> None: + """ + Build one saved entry with the create the API uses for its kind, and register it under its name. + + The entry passes the same checks a create request does. Restore runs as the server, and a saved file + doesn't record who saved it, so a target is built as an administrator would build it. A target that + uploads files gets this server's upload folder, as it does through the API. + + Args: + stored (StoredRecipe): The checked recipe to build. + upload_directory (str | None): The folder this server lets targets upload files from, or None. + + Raises: + RecipeRejectedError: If the recipe doesn't make a valid create request. + """ + request = _create_request(stored) + if isinstance(request, CreateTechniqueRequest): + await get_technique_service().create_async(request) + elif isinstance(request, CreateTargetRequest): + await get_target_service().create_target_async( + request=request, upload_directory=upload_directory, is_admin=True + ) + elif isinstance(request, CreateConverterRequest): + await get_converter_service().create_converter_async(request=request) + else: + await get_scorer_service().create_scorer_async(request=request) + + +def _create_request( + stored: StoredRecipe, +) -> CreateTargetRequest | CreateConverterRequest | CreateScorerRequest | CreateTechniqueRequest: + """ + Turn a saved recipe back into the create request it was saved from. + + Returns: + CreateTargetRequest | CreateConverterRequest | CreateScorerRequest | CreateTechniqueRequest: The request. + + Raises: + RecipeRejectedError: If the recipe doesn't make a valid request, such as a technique with a reserved name. + """ + recipe = stored.recipe + fields: dict[str, Any] = {"name": stored.ref.name} + try: + if isinstance(recipe, TechniqueRecipe): + return CreateTechniqueRequest.model_validate( + fields | recipe.model_dump(exclude_unset=True, exclude={"schema_version"}) + ) + fields |= {"type": recipe.class_name, "params": recipe.params} + if stored.ref.kind is RegistryRecipeKind.TARGETS: + if recipe.auth_mode is not None: + fields["auth_mode"] = recipe.auth_mode + return CreateTargetRequest.model_validate(fields) + if stored.ref.kind is RegistryRecipeKind.CONVERTERS: + return CreateConverterRequest.model_validate(fields) + return CreateScorerRequest.model_validate(fields) + except ValidationError as error: + problems = "; ".join( + f"{'.'.join(map(str, issue['loc'])) or 'request'}: {issue['msg']}" for issue in error.errors() + ) + raise RecipeRejectedError( + reason=RestoreFailureReason.INVALID_RECIPE, message=f"Not a valid create request: {problems}" + ) from error diff --git a/pyrit/backend/services/scenario_suite_service.py b/pyrit/backend/services/scenario_suite_service.py index 65cf984e94..8fbe7546a5 100644 --- a/pyrit/backend/services/scenario_suite_service.py +++ b/pyrit/backend/services/scenario_suite_service.py @@ -5,17 +5,17 @@ import asyncio import logging +import traceback from dataclasses import dataclass from functools import lru_cache from typing import Literal -from azure.core.exceptions import AzureError - from pyrit.backend.models import ( ScenarioSuiteEntryStatus, ScenarioSuiteListResponse, ScenarioSuiteResponse, ) +from pyrit.common.azure_storage import redact_urls from pyrit.models import ( PresetIssue, ScenarioSuite, @@ -25,6 +25,7 @@ ScenarioSuiteRegistryEntry, ) from pyrit.registry import ScenarioPresetRegistry, ScenarioSuiteRegistry +from pyrit.registry.file_document_storage import DocumentStorageUnavailableError logger = logging.getLogger(__name__) @@ -217,12 +218,10 @@ def _resolve_preset_projection( ) stored_version = registry.get_preset_version(name=preset_name) - except (AzureError, OSError): - logger.warning( - "Could not project scenario preset '%s' because preset storage is unavailable.", - preset_name, - exc_info=True, - ) + except DocumentStorageUnavailableError as error: + # The error behind it can carry a SAS URL, so the traceback is logged with URL credentials removed. + details = redact_urls("".join(traceback.format_exception(error)).rstrip()) + logger.warning(f"Could not project scenario preset '{preset_name}': {error}\n{details}") return _PresetProjection( origin=None, version=None, diff --git a/pyrit/backend/services/target_service.py b/pyrit/backend/services/target_service.py index a3d413e411..91c0205e61 100644 --- a/pyrit/backend/services/target_service.py +++ b/pyrit/backend/services/target_service.py @@ -39,6 +39,7 @@ from pyrit.models.component_spec import SourceInstanceSpec, TargetBinding from pyrit.prompt_target import ( GitHubCopilotTarget, + HuggingFaceChatTarget, PromptTarget, RoundRobinTarget, discover_target_capabilities_async, @@ -48,7 +49,7 @@ CapabilityDiscoveryError, ) from pyrit.prompt_target.common.tool_call_history import TOOL_CALL_INPUT_MODALITIES -from pyrit.registry import TargetRegistry +from pyrit.registry import TargetMetadata, TargetRegistry logger = logging.getLogger(__name__) @@ -65,6 +66,24 @@ _AUTH_MODE_PARAM = "auth_mode" +def _import_torch(*, use_cuda: bool) -> None: + """ + Import torch, and check for CUDA when asked, as a HuggingFaceChatTarget constructor does. + + Either can take seconds the first time and is quick after it, so a worker thread does them before the + constructor runs on the event loop. A missing torch is left for the constructor to report. + + Args: + use_cuda (bool): Whether the target asks for CUDA. + """ + try: + import torch # type: ignore[ty:unresolved-import] + except ModuleNotFoundError: + return + if use_cuda: + torch.cuda.is_available() + + class TargetService: """ Service for managing target instances. @@ -651,6 +670,11 @@ async def create_target_async( is not target_cls ): raise ValueError("The source target type does not match the requested type") + if issubclass(target_cls, HuggingFaceChatTarget): + # Its constructor imports torch and checks for CUDA, which can hold this loop for seconds the first time, + # then starts loading the model on this loop. Doing both first in a worker leaves the constructor quick. + use_cuda = self._requests_cuda(metadata=metadata, params=permission_params) + await asyncio.to_thread(_import_torch, use_cuda=use_cuda) target_obj = ( await self.build_from_source_async(request.source) if request.source @@ -675,6 +699,29 @@ async def create_target_async( await release_component_async(target_obj) raise + @staticmethod + def _requests_cuda(*, metadata: TargetMetadata | None, params: dict[str, Any]) -> bool: + """ + Read ``use_cuda`` as the constructor will, through its registered parameter. + + A textual boolean the catalog accepts, such as "false", is read the same way. A value the registry can't read + as a boolean counts as false here, since construction rejects it. + + Args: + metadata (TargetMetadata | None): The target class's registered metadata. + params (dict[str, Any]): The constructor parameters. + + Returns: + bool: Whether the target asks for CUDA. + """ + parameter = next((p for p in (metadata.parameters if metadata else ()) if p.name == "use_cuda"), None) + if parameter is None: + return False + try: + return bool(parameter.coerce_value(params.get("use_cuda", parameter.default))) + except ValueError: + return False + async def _enforce_code_execution_permission_async( self, *, target_type: str, params: dict[str, Any], is_admin: bool ) -> None: diff --git a/pyrit/common/azure_storage.py b/pyrit/common/azure_storage.py index d3ed6eeb4f..dfce6b2e08 100644 --- a/pyrit/common/azure_storage.py +++ b/pyrit/common/azure_storage.py @@ -3,8 +3,16 @@ """Lightweight Azure Storage URI helpers.""" +import re from urllib.parse import parse_qs, urlparse +# A scheme followed by "://" and everything up to whitespace, so a quote can't end a URL before its query. +# A scheme starts only where one can't continue from the character before it, such as after "_" or a +# space, so each run of scheme characters is scanned once and long text stays fast. +_URL_IN_TEXT = re.compile(r"(? bool: """ @@ -48,9 +56,30 @@ def has_sas_signature(value: str) -> bool: def redact_url_credentials(value: str) -> str: """ - Remove query parameters and fragments from a URL for display or logging. + Remove the user name and password, query parameters, and fragment from a URL for display or logging. + + Returns: + str: The URL without them. + """ + parsed = urlparse(value) + return parsed._replace(netloc=parsed.netloc.rpartition("@")[2], query="", fragment="").geturl() + + +def redact_urls(text: str) -> str: + """ + Remove the user name and password, query parameters, and fragment from every URL in a text. + + A query or user name can hold a credential, such as a SAS token. A URL that can't be parsed is + hidden entirely, and a SAS signature is hidden wherever it appears. Returns: - str: The URL without query parameters or a fragment. + str: The text with each URL redacted. """ - return urlparse(value)._replace(query="", fragment="").geturl() + + def redact(match: re.Match[str]) -> str: + try: + return redact_url_credentials(match.group(0)) + except ValueError: + return "