Skip to content

feat(library): add Agent Threat Rules (ATR) detection rail - #2251

Open
eeee2345 wants to merge 1 commit into
NVIDIA-NeMo:developfrom
eeee2345:feat/atr-rail-manifest
Open

feat(library): add Agent Threat Rules (ATR) detection rail#2251
eeee2345 wants to merge 1 commit into
NVIDIA-NeMo:developfrom
eeee2345:feat/atr-rail-manifest

Conversation

@eeee2345

@eeee2345 eeee2345 commented Aug 4, 2026

Copy link
Copy Markdown

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 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 yields the same verdict on every run. It is modelled on injection_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.py fails test_library_action_packages_declare_manifests. Rebasing was not enough; the module needed rewriting to the manifest pattern.

Areas worth careful review:

  • rail.py declares RailPrivacy() 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.py returns RailOutcome and 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.json is regenerated. The diff is purely additive (+25 / -0) and contains only the atr config key and ATRDetection definition.

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 flow rather than define 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 as injection_detection_reject_blocks_injection failing with an ATR refusal string in place of the injection one. Both flows now emit the shared bot refuse to respond intent instead of an inline templated bot say, and test_colang_1_flow_is_declared_as_a_subflow guards 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, pyatr 0.2.7:

  • tests/test_atr_rail.py, tests/manifests/, tests/rails/llm/, tests/test_runtime_flow_gate_equivalence.py399 passed.
  • Baselined every failure against clean develop before treating it as mine; all four initial failures were caused by this change and are fixed, not suppressed.
  • ruff format --check and ruff check clean on the touched files.
  • uv lock --check passes.

Not run: the recorded/live-credential suites and the docs build.

Residual risk: uv.lock is regenerated with uv 0.10.1. The marker normalization and exclude-newer churn in that file comes from the uv version rather than from this change — uv-latest.yml already runs uv lock --check on latest uv, but if you pin a different version please regenerate.

AI Assistance

  • No AI tools were used.
  • AI tools were used; a human reviewed and can explain every change (tool: Claude Code).

Checklist

  • I've read the CONTRIBUTING guidelines.
  • This PR links to a triaged issue assigned to me. — Add an Agent Threat Rules (ATR) detection library rail #1991 is open but not yet triaged or assigned; requested in the issue.
  • My PR title follows the project commit convention.
  • I've updated the documentation if applicable.
  • I've added tests if applicable.
  • I've noted any verification beyond CI and any checks I couldn't run.
  • I did not update generated changelog files manually.
  • I addressed all CodeRabbit, Greptile, and other review comments, or replied with why no change is needed. — none yet on this PR.
  • @mentions of the person or team responsible for reviewing proposed changes. — @Pouyanpi

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.
@github-actions github-actions Bot added status: needs triage New issues that have not yet been reviewed or categorized. size: L labels Aug 4, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

size: L status: needs triage New issues that have not yet been reviewed or categorized.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant