FEAT Wire frontend attack view to backend APIs#1371
FEAT Wire frontend attack view to backend APIs#1371romanlutz wants to merge 1 commit intoAzure:mainfrom
Conversation
Add target management and real attack execution to the PyRIT frontend, replacing the echo stub with live backend communication. Backend: - Add attack CRUD endpoints (create, list, get, messages, add message) - Replace target_unique_name with target_registry_name across models, routes, services, and mappers for clarity - Add TargetInfo model to nest target metadata in AttackSummary - Add target_registry_name to AddMessageRequest so the backend stays stateless (no reverse lookups from attack identifier) - Move lifespan init to pyrit_backend CLI; main.py warns if run directly CLI: - Add FrontendCore.run_initializers_async() to consolidate initializer resolution used by both pyrit_backend and run_scenario_async - Move all deferred imports to module level in frontend_core.py - Refactor pyrit_backend to use FrontendCore two-step init pattern (initialize_async + run_initializers_async) Frontend: - Add Config page with target list, create target dialog, and set-active-target flow - Wire ChatWindow to attacksApi: lazy attack creation on first message, send via PromptNormalizer, map backend responses to UI messages - Add messageMapper utils (backend DTO <-> frontend Message conversion) - Add full backend DTO types mirroring pyrit/backend/models - Support simulated_assistant role, error rendering, loading indicators, and media attachments (image/audio/video) Tests: - Add/update 300+ backend unit tests covering attack service, mappers, target service, API routes, and main lifespan - Add 150+ frontend tests covering ChatWindow, TargetConfig, CreateTargetDialog, MessageList, api service, and messageMapper - Update test_frontend_core patch targets to match top-level imports
| # Initialization is handled by the pyrit_backend CLI before uvicorn starts. | ||
| # Running 'uvicorn pyrit.backend.main:app' directly is not supported; | ||
| # use 'pyrit_backend' instead. | ||
| if not CentralMemory._memory_instance: |
There was a problem hiding this comment.
We should probably not access private data-member in CentralMemory. That breaks encapsulation. Maybe we could do a try/except block around get_memory_instance() or add a public method to check the initialization, e.g. is_initialized()
| async def run_initializers_async(self) -> None: | ||
| """ | ||
| Resolve and run all configured initializers and initialization scripts. | ||
|
|
||
| Must be called after :meth:`initialize_async` so that registries are | ||
| available to resolve initializer names. This is the same pattern used | ||
| by :func:`run_scenario_async` before executing a scenario. | ||
|
|
||
| If no initializers are configured this is a no-op. | ||
| """ | ||
| initializer_instances = None | ||
| if self._initializer_names: | ||
| print(f"Running {len(self._initializer_names)} initializer(s)...") | ||
| sys.stdout.flush() | ||
| initializer_instances = [self.initializer_registry.get_class(name)() for name in self._initializer_names] | ||
|
|
||
| await initialize_pyrit_async( | ||
| memory_db_type=self._database, | ||
| initialization_scripts=self._initialization_scripts, | ||
| initializers=initializer_instances, | ||
| env_files=self._env_files, | ||
| ) | ||
|
|
There was a problem hiding this comment.
I see your concern here, and it is very valid. Looking at the code, the issue here is that initialize_pyrit_async is a big function that bundles 4 different responsibilities:
- Load env files
- Reset default values
- Create and set memory
- Run initializers and scripts
You need to do 1-3 once at startup, then 4 separately (or per scenario). But since there's no way to run just 4, you call the entire function twice, basically recreating memory, reresetting defaults, and reloading env files.
I think the easiest way to address this is to extract a new top-level function and then refactor initialize_pyrit_async to use that. For example, in the initialization file, you could add this:
async def run_initializers_async(
*,
initialization_scripts: Optional[Sequence[Union[str, pathlib.Path]]] = None,
initializers: Optional[Sequence["PyRITInitializer"]] = None,
) -> None:
# this is just to make sure memory is initialized, it will raise an exception if it's not
CentralMemory.get_memory_instance()
reset_default_values()
all_initializers = list(initializers) if initializers else []
if initialization_scripts:
script_initializers = _load_initializers_from_scripts(script_paths=initialization_scripts)
all_initializers.extend(script_initializers)
if all_initializers:
await _execute_initializers_async(initializers=all_initializers)This would be a small change, and it is also backward compatible so no need to update the callsites.
Then in this file you could do this:
from pyrit.setup.initialization import run_initializers_async as _run_initializers_async
async def initialize_async(self) -> None:
if self._initialized:
return
await initialize_pyrit_async(
memory_db_type=self._database,
initialization_scripts=None,
initializers=None,
env_files=self._env_files,
)
self._scenario_registry = ScenarioRegistry.get_registry_singleton()
if self._initialization_scripts:
print("Discovering user scenarios...")
sys.stdout.flush()
self._scenario_registry.discover_user_scenarios()
self._initializer_registry = InitializerRegistry()
self._initialized = True
async def run_initializers_async(self) -> None:
initializer_instances = None
if self._initializer_names:
print(f"Running {len(self._initializer_names)} initializer(s)...")
sys.stdout.flush()
initializer_instances = [self.initializer_registry.get_class(name)() for name in self._initializer_names]
# this will now use the new function you created, no duplicate work
await _run_initializers_async(
initialization_scripts=self._initialization_scripts,
initializers=initializer_instances,
)You could also refactor initialize_pyrit_async to use the new function:
async def initialize_pyrit_async(...) -> None:
_load_environment_files(env_files=env_files, silent=silent)
reset_default_values()
# could we move all the memory initialization to a util function :)?
_setup_memory(memory_db_type, **memory_instance_kwargs)
if initializers or initialization_scripts:
await run_initializers_async(
initialization_scripts=initialization_scripts,
initializers=initializers,
)| @@ -91,9 +96,3 @@ def setup_frontend() -> None: | |||
|
|
|||
| # Set up frontend at module load time (needed when running via uvicorn) | |||
| setup_frontend() | |||
There was a problem hiding this comment.
maybe this should go into lifespan?
| const [targetType, setTargetType] = useState<string>('') | ||
| const [endpoint, setEndpoint] = useState('') | ||
| const [modelName, setModelName] = useState('') | ||
| const [apiKey, setApiKey] = useState('') | ||
| const [submitting, setSubmitting] = useState(false) |
There was a problem hiding this comment.
why are you explicitly typing in some of these
|
|
||
| const handleSubmit = async () => { | ||
| if (!targetType) { | ||
| setError('Please select a target type') |
There was a problem hiding this comment.
you can set this on the field itself: https://storybooks.fluentui.dev/react/?path=/docs/components-field--docs
| import CreateTargetDialog from './CreateTargetDialog' | ||
|
|
||
| const useStyles = makeStyles({ | ||
| root: { |
There was a problem hiding this comment.
we should probably have a separate styles file
| // Wait before retrying (1s, 2s, 3s) | ||
| await new Promise(r => setTimeout(r, (attempt + 1) * 1000)) | ||
| } else { | ||
| setError(err instanceof Error ? err.message : 'Failed to load targets') |
There was a problem hiding this comment.
when isn't it going to be an instance of an error ?
|
|
||
| {!loading && !error && targets.length > 0 && ( | ||
| <div className={styles.tableContainer}> | ||
| <Table aria-label="Target instances"> |
There was a problem hiding this comment.
I think we should separate this table into a different component
|
Can you add screenshots for the frontend |
| // returns 502 while the backend is still starting, so a single failed | ||
| // request on initial page load would show a confusing error to the user. | ||
| const fetchTargets = useCallback(async () => { | ||
| const maxRetries = 3 |
There was a problem hiding this comment.
I'm not sure I understand this logic for retries. The comment mentions that the dev proxy returns 502 while the backend is still starting. Doesn't that mean we know the cause of failure? Why not return something like "backend still loading" to the user and treat it as a different error path?
| <TableHeaderCell>Model</TableHeaderCell> | ||
| </TableRow> | ||
| </TableHeader> | ||
| <TableBody> |
There was a problem hiding this comment.
The table entries (anything under <TableCell>) will not auto-truncate to my knowledge. Do you think we should add something in the CSS (or change these elements) to enforce a character limit?
Add target management and real attack execution to the PyRIT frontend, replacing the echo stub with live backend communication.
Backend:
CLI:
Frontend:
Tests: