FEAT Plug-in mechanism for private scenarios

[ValbuenaVC]

Description

Operators who use PyRIT extensively sometimes need to run bespoke scenario logic that cannot be published to the public repo and cannot be represented as database rows. Today there is no delivery path for that code into a stock public install. This PR adds a minimal plug-in mechanism so a private scenario can be shipped as a pre-built Python wheel and loaded at initialization, without those components living in the public repository. Datasets come along as a secondary benefit (a plug-in's dataset providers self-register on import), but the primary purpose is distributing private scenario code.

TL; DR How it works. Plug-in loading runs as a guaranteed-first phase inside initialize_pyrit_async, after central memory is set and before the configured initializers run. Because it is a fixed phase rather than a user-ordered .pyrit_conf initializer, "plug-ins register before anything reads the registry" is true by construction. This deliberately avoids putting a load-bearing ordering requirement into user-editable config (a silent-failure trap: a misordered conf would let scenario metadata warm before plug-in scenarios registered). Set PLUGIN_WHEEL to a wheel on disk and the loader extracts it (stdlib zipfile, never pip or .venv) into .plugin/<name>/, prepends that dir to sys.path, imports the package (dataset providers self-register), runs the plug-in's bootstrap (a top-level register() callable or a shipped PyRITInitializer subclass that registers scenarios), and asserts the plug-in registered something. It is a no-op when PLUGIN_WHEEL is unset.

Optional config: PLUGIN_PACKAGE (auto-detected top-level package), PLUGIN_FAIL_OPEN (or initialize_pyrit_async(plugin_fail_open=True)), and PLUGIN_DIR. It fails closed by default; fail-open skips a broken plug-in with a warning.

Why extraction, not install. Plug-in seed files load via __file__-relative paths; a raw .whl on sys.path (zipimport) makes those resolve inside the zip and the datasets silently vanish. Runtime pip/setuptools are also not guaranteed present (uv venvs ship without them). A wheel is just a zip, so stdlib extraction works on any environment.

Silent-failure guards (the non-obvious part). Each guard exists because it is a way a plug-in could load "successfully" yet be wrong: extraction rather than zipimport; atomic extraction (temp dir then move) so a crash never leaves a poisoned cache; submodule discovery so providers and the bootstrap register even if the package __init__ does not import them; a package-shadowing guard that fails loudly if an installed package of the same name shadows the wheel; fail-loud when a wheel imports cleanly but registers nothing; and full rollback of sys.path, imported modules, and registries on any failure, so a failed (or fail-open) load leaves no trace.

Dataset guardrail (worth a careful look). The dataset resolver treats central memory as authoritative and only consults a provider when memory has no seeds for that name, so a plug-in cannot override an existing same-named dataset (a scan would silently use the existing copy). To surface this, the loader warns loudly at load time (greppable PLUGIN DATASET SHADOWED: prefix, naming both providers) when a plug-in dataset_name collides with an existing registered provider's name. The check compares the provider registry rather than live memory on purpose: at that phase memory is not populated yet, and a live-memory check would false-positive on the plug-in's own datasets persisted from a prior run (the source is undecidable at load time). Precedence for names that exist only in shared storage without a provider is out of scope here and belongs to a separate dataset-management effort.

Trust boundary. Running a plug-in bootstrap executes third-party code in the PyRIT process, so whoever can set PLUGIN_* or write the backend's .env can run code on the host. The config is the trust boundary and should be treated as sensitive.

Scoped out of this PR (natural follow-ups): custom attacks/converters/targets as plug-ins, a live-iteration/editable-install workflow, nested dependency resolution, multi-plug-in composition, and .pyrit_conf composition.

Tests and Documentation

Tests: tests/unit/setup/test_plugin_loader.py builds a mock plug-in wheel at test time (no dependency on any real plug-in) and covers the full spine plus every guard: extract-not-zipimport, provider self-registration on import (including via submodule discovery), scenario register_class surviving lazy discovery and coexisting with built-ins, the no-arg-instantiable contract, the phase running after memory is set and before the configured initializers, package-name resolution, atomic-extraction cache reuse, the package-shadowing guard, the dataset-name collision warning (and no false positive on the plug-in's own reload), and all failure modes (fail-closed, fail-open, loud-on-empty, rollback). 43 test cases; the full tests/unit/setup/ suite passes (323). ruff, ty, and the async-suffix check are clean.

Documentation: operator-facing configuration is documented inline in .env_example (the PLUGIN_* variables) and .pyrit_conf_example (a note that plug-ins are not configured there). No Jupyter notebooks changed, so JupyText was not run.

[Copilot]

Pull request overview

Adds a wheel-based plug-in loading mechanism to PyRIT initialization so operators can ship private scenarios (and dataset providers) outside the public repo, while ensuring plug-ins load before other initializers consume registries.

Changes:

• Introduces pyrit.setup.plugin_loader to extract a configured .whl into .plugin/, import it, run a bootstrap (register() or PyRITInitializer), and validate registration with rollback-on-failure.
• Hooks plug-in loading into initialize_pyrit_async as a guaranteed-first phase (after memory is set, before other initializers), with a plugin_fail_open override.
• Adds extensive unit coverage for loader behavior/guards and documents operator configuration via .env_example and .pyrit_conf_example; ensures .plugin/ is kept but ignored in git.

Reviewed changes

Copilot reviewed 6 out of 7 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
tests/unit/setup/test_plugin_loader.py	New unit test suite that builds a mock wheel and exercises loader spine + guardrails (fail-open/closed, rollback, collisions, cache, package inference).
pyrit/setup/plugin_loader.py	New plug-in loader implementation: extraction, import/submodule discovery, bootstrap execution, collision warnings, and rollback.
pyrit/setup/initialization.py	Adds the plug-in load phase into initialize_pyrit_async and exposes plugin_fail_open.
.pyrit_conf_example	Documents that plug-ins are not configured via .pyrit_conf and are loaded automatically during initialization.
.plugin/.gitkeep	Keeps .plugin/ directory present in the repo while allowing extracted artifacts to be ignored.
.gitignore	Ignores extracted plug-in artifacts under .plugin/ while retaining .gitkeep.
.env_example	Documents PLUGIN_WHEEL, PLUGIN_PACKAGE, PLUGIN_FAIL_OPEN, and PLUGIN_DIR configuration and trust boundary.

[fdubut]

Also as mentioned in our side conversation, I think it would be useful to share an example of how a customer would package their datasets/scenarios to make them available as a plugin.

[fdubut]

Can we find a naming for this flag that's a bit clearer about what it does?

[ValbuenaVC]

Yes. Will patch this

[fdubut]

Is .env the right place for this? Intuitively I would have put it in .pyrit.conf.

[ValbuenaVC]

That's a much better place for it! I agree, I'll move it

[fdubut]

Does that mean we support only one plugin? If so, I think it's a design limitation that might bite us later.

[ValbuenaVC]

To answer your first question, yes, but not as an explicit constraint. I see two solutions to plug-in composition: we can just load multiple plug-ins and place the burden of compatibility on the user, or we can add a layer of abstraction, so that plug-in components can be collected, compared, and initialized in a certain order. I lean towards the second long-term, but in either case I don't think we should support only one plugin on purpose.

[fdubut]

I would even generalize the use case, it's not just about private datasets and scenarios from AIRT but any customer may want to package their code in a plugin that doesn't live collocated with the main PyRIT code.

[ValbuenaVC]

For this PR specifically I'm limiting test coverage to scenarios and datasets for tractability, but this is the long-term intent of this feature.

I want to publish a follow-up PR that handles building the plug-ins for users. I imagine it could look something like this:

pyrit_scan --build-plugin C:\Users\roakey\Dev\operation-super-secret\ --install-after y
Building plug-in from directory \`operation-super-secret\`...
[INFO] Installing plug-in \`A\` with initializer \`InitializerA\`...
[INFO] Discovered 3 scenarios...
[INFO] Scenarios all instantiated with no problems!
[INFO] Installing plug-in \`B\` with initializer \`InitializerB\`...
[INFO] Discovered 12 datasets...
[WARN] Dataset Foo is associated with file \`bar.prompt\`, but bar.prompt was not found! Foo has been excluded.
Done!

The plugin builder solves the packaging problem. A related problem is component support. There is nothing in principle stopping users from extending this to every PyRIT component. But guaranteeing this works smoothly presents a number of challenges I've chosen to scope out of this PR. Among them: naming collisions, syncing dependencies, race conditions/initializer precedence, pyrit/plugin version drift.