feat(library): add Agent Threat Rules (ATR) detection rail - #2251
Open
eeee2345 wants to merge 1 commit into
Open
feat(library): add Agent Threat Rules (ATR) detection rail#2251eeee2345 wants to merge 1 commit into
eeee2345 wants to merge 1 commit into
Conversation
Adds an input rail that matches the user message against Agent Threat Rules, an open MIT-licensed detection catalog for AI-agent attacks (prompt injection, jailbreak, tool poisoning, MCP attacks, skill compromise). Rules ship inside the optional `pyatr` package and are evaluated in-process, so the rail needs no API key, no service endpoint and no network access, and the same input produces the same verdict on every run. This supersedes NVIDIA-NeMo#1992, which predates the rail-manifest migration and could not have merged: a library package with actions and no `rail.py` fails `test_library_action_packages_declare_manifests`. - `rail.py` / `rail_config.py`: manifest and typed config section, following `injection_detection` (the other local, optional-dependency rail). The manifest declares `optional_dependencies=("pyatr",)` and a `RailPrivacy` with every disclosure flag false and no remote services. - `actions.py`: returns `RailOutcome` so the surface is portable, per `test_portable_surface_actions_declare_rail_outcome_returns`. Matched rule IDs and the max severity ride along in metadata, including on allow, so sub-threshold matches still show up in traces. - Colang 1 flow is a `define subflow`, not a `define flow`. As a plain flow it became a top-level flow the runtime activated on its own: it ran for configs that never asked for it and overrode the configured rail's refusal message, rendering an ATR message with an empty rule list. That is what broke `injection_detection_reject_blocks_injection` locally. - Both flows emit the shared `bot refuse to respond` intent instead of an inline templated `bot say`, so refusal wording stays configurable in one place. - `block_severities: []` is honored as monitor-only; only an absent config section falls back to `["critical", "high"]`. Tests: rail behaviour, an `atr` extra, flow-gate equivalence cases via `_rail_outcome_cases`, and a structural guard on the subflow declaration. 399 tests pass across tests/test_atr_rail.py, tests/manifests, tests/rails/llm and tests/test_runtime_flow_gate_equivalence.py. uv.lock is regenerated; the marker and exclude-newer churn comes from the uv version, not from this change.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Adds an input rail that matches the user message against Agent Threat Rules, an open MIT-licensed detection catalog for AI-agent attacks (prompt injection, jailbreak, tool poisoning, MCP attacks, skill compromise). Rules ship inside the optional
pyatrpackage and are evaluated in-process, so the rail needs no API key, no service endpoint and no network access, and the same input yields the same verdict on every run. It is modelled oninjection_detection, the other local optional-dependency rail.This supersedes #1992, which predates the rail-manifest migration and could not have merged in its shape — a library package with actions and no
rail.pyfailstest_library_action_packages_declare_manifests. Rebasing was not enough; the module needed rewriting to the manifest pattern.Areas worth careful review:
rail.pydeclaresRailPrivacy()with every disclosure flag false and no remote services. Please sanity-check that this is the intended way to express a rail that contacts nothing, rather than simply leaving the field at its default.actions.pyreturnsRailOutcomeand carries matched rule IDs plus max severity in metadata on allow as well as block, so sub-threshold matches remain visible in traces. Say the word if you would rather allow paths carried empty metadata.schemas/rails_config.snapshot.jsonis regenerated. The diff is purely additive (+25 / -0) and contains only theatrconfig key andATRDetectiondefinition.One defect found while porting, flagged separately because it is the kind of thing that would have been unpleasant in a release: the Colang 1 flow in #1992 was declared
define flowrather thandefine subflow. As a top-level flow the runtime activated it independently of configuration — it ran for configs that never enabled the rail and overrode the configured rail's refusal message, rendering an ATR message with an empty rule list. Locally that showed up asinjection_detection_reject_blocks_injectionfailing with an ATR refusal string in place of the injection one. Both flows now emit the sharedbot refuse to respondintent instead of an inline templatedbot say, andtest_colang_1_flow_is_declared_as_a_subflowguards the declaration.block_severities: []is honored as monitor-only; only an absent config section falls back to["critical", "high"].Related Issue(s)
Verification
Run locally against
develop@origin/develop, Python 3.12,pyatr0.2.7:tests/test_atr_rail.py,tests/manifests/,tests/rails/llm/,tests/test_runtime_flow_gate_equivalence.py— 399 passed.developbefore treating it as mine; all four initial failures were caused by this change and are fixed, not suppressed.ruff format --checkandruff checkclean on the touched files.uv lock --checkpasses.Not run: the recorded/live-credential suites and the docs build.
Residual risk:
uv.lockis regenerated with uv 0.10.1. The marker normalization andexclude-newerchurn in that file comes from the uv version rather than from this change —uv-latest.ymlalready runsuv lock --checkon latest uv, but if you pin a different version please regenerate.AI Assistance
Checklist