diff --git a/python/samples/02-agents/middleware/README.md b/python/samples/02-agents/middleware/README.md index 4b7d0bbdfb0..99ad31a454c 100644 --- a/python/samples/02-agents/middleware/README.md +++ b/python/samples/02-agents/middleware/README.md @@ -11,7 +11,7 @@ This folder contains focused middleware samples for `Agent`, chat clients, tools | [`agent_loop_middleware_todos.py`](./agent_loop_middleware_todos.py) | Demonstrates `AgentLoopMiddleware` with a `should_continue` predicate built from a `TodoProvider` via `todos_remaining`, so the agent keeps working while open todos remain. | | [`agent_loop_middleware_judge.py`](./agent_loop_middleware_judge.py) | Demonstrates `AgentLoopMiddleware.with_judge`: a ChatClient judge re-runs the agent until it decides the original request was answered, with `criteria` shared between the agent and the judge. | | [`agent_loop_middleware_report.py`](./agent_loop_middleware_report.py) | Demonstrates composing two `AgentLoopMiddleware` on one agent: an inner `todos_remaining` loop that drafts a report todo-by-todo, wrapped by an outer report-style `with_judge` loop that re-runs it until an editor chat client judges the report publication-ready. | -| [`atr_validation_middleware.py`](./atr_validation_middleware.py) | Demonstrates deterministic validation at the tool-execution boundary: a `FunctionMiddleware` that inspects the validated tool arguments and raises `MiddlewareTermination` before the tool runs when they match an attack rule. Loads the open, MIT-licensed Agent Threat Rules ruleset and runs the real engine locally (`pip install pyatr`), with a built-in deny-list fallback when it is not installed. | +| [`atr_validation_middleware.py`](./atr_validation_middleware.py) | Demonstrates deterministic validation at the tool-execution boundary: a `FunctionMiddleware` that raises `MiddlewareTermination` before the tool runs when the arguments match a high- or critical-severity Agent Threat Rules rule, and logs lower-severity matches as warnings. It is not a dangerous-command filter, so use tool approval in production. Runs the open, MIT-licensed ruleset locally (`pip install pyatr`). | | [`chat_middleware.py`](./chat_middleware.py) | Shows class-based and function-based chat middleware that can observe, modify, and override model calls. | | [`class_based_middleware.py`](./class_based_middleware.py) | Shows class-based agent and function middleware. | | [`decorator_middleware.py`](./decorator_middleware.py) | Demonstrates middleware registration with decorators. | diff --git a/python/samples/02-agents/middleware/atr_validation_middleware.py b/python/samples/02-agents/middleware/atr_validation_middleware.py index 1245439e3a6..eaa2495b94f 100644 --- a/python/samples/02-agents/middleware/atr_validation_middleware.py +++ b/python/samples/02-agents/middleware/atr_validation_middleware.py @@ -36,7 +36,22 @@ point that validates a tool call right before it executes. ATRValidationMiddleware is a FunctionMiddleware that inspects the validated tool arguments in ``FunctionInvocationContext.arguments`` and raises ``MiddlewareTermination`` BEFORE calling -``call_next()`` when the arguments match a known attack pattern, so the tool never runs. +``call_next()`` when the arguments match a high- or critical-severity attack rule, so the tool +never runs. + +Why that threshold: ATR's reference verdict is "deny" or "ask" for every high- or +critical-severity match, where "ask" means a person should confirm the call. This sample has no +confirmation step (``approval_mode="never_require"``), so it treats both as stop. Below high, the +cut is a false-positive trade-off, not ATR's verdict: ATR also returns "ask" for some +medium-severity matches, but those rules also match ordinary commands such as +``curl https://example.com``, so medium- and low-severity matches are logged as warnings and the +tool runs. + +Limits: ATR targets agent attacks such as prompt injection, tool-argument tampering, and +exfiltration; it is not a filter for dangerous shell commands. ``rm -rf ~`` matches only a +medium-severity rule, so it runs with a warning, and some destructive commands match no rule at +all. This middleware does not replace tool approval (``approval_mode="always_require"``) for +tools with side effects. Detection is delegated to Agent Threat Rules (ATR) -- an open, MIT-licensed detection ruleset for AI-agent threats such as prompt injection, tool-argument tampering, and exfiltration. The @@ -48,6 +63,10 @@ logger = logging.getLogger(__name__) +# ATR's verdict for a high or critical match is deny or ask; with no confirmation step here, both +# terminate the run. Medium and low matches are logged and the tool runs. +BLOCKING_SEVERITIES = frozenset({"critical", "high"}) + def _arguments_to_text(arguments: BaseModel | Mapping[str, Any]) -> str: """Flatten tool arguments into a single string for scanning. @@ -71,27 +90,29 @@ def _load_atr_engine() -> Any: return engine -def detect_attack(arguments: BaseModel | Mapping[str, Any]) -> str | None: - """Return the matched ATR rule id, or None when the arguments look benign. +def detect_attack(arguments: BaseModel | Mapping[str, Any]) -> Any | None: + """Return the highest-severity ATR match, or None when the arguments look benign. Runs the real ATR engine over the flattened tool arguments. The text is evaluated as a ``tool_call`` event so it is checked against the rules' ``tool_args`` conditions; ``evaluate`` - sorts matches critical-first, so the first rule id is the highest-severity hit. + sorts matches critical-first, so the first match is the highest-severity hit. The returned + ``pyatr.ATRMatch`` carries ``rule_id``, ``severity`` and ``confidence``; the middleware decides + whether to block. The ruleset replaces a hand-rolled deny-list. For reference, the shape of the patterns ATR - encodes (and that the earlier version of this sample inlined) is, e.g.:: + encodes is, e.g.:: ignore (previous|prior|above) instructions # instruction override / prompt injection send (secret|token|api_key|password) to http... # credential exfiltration (cat|read|open) (.env|id_rsa|/etc/passwd) # sensitive-file access - pyatr ships hundreds of such rules and keeps them maintained, so the sample stays a single - straight-line call instead of a local regex list. + pyatr ships and maintains these rules, so the sample stays a single straight-line call + instead of a local regex list. """ text = _arguments_to_text(arguments) event = pyatr.AgentEvent(content=text, event_type="tool_call", fields={"tool_args": text}) matches = _load_atr_engine().evaluate(event) - return matches[0].rule_id if matches else None + return matches[0] if matches else None # NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production; @@ -107,10 +128,12 @@ def get_weather( class ATRValidationMiddleware(FunctionMiddleware): - """Validates tool arguments at the execution boundary and blocks malicious calls. + """Validates tool arguments at the execution boundary and blocks high- or critical-severity matches. - The check is deterministic and runs before the tool executes: on a match it raises + The check is deterministic and runs before the tool executes: on a match whose severity is in + ``BLOCKING_SEVERITIES`` (ATR verdict deny or ask, with no one here to confirm) it raises ``MiddlewareTermination`` so ``call_next()`` is never reached and the tool does not fire. + Medium- and low-severity matches are logged as warnings and the tool runs. """ async def process( @@ -118,18 +141,32 @@ async def process( context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]], ) -> None: - matched = detect_attack(context.arguments) - if matched is not None: + match = detect_attack(context.arguments) + if match is not None and match.severity in BLOCKING_SEVERITIES: logger.warning( - "[ATRValidationMiddleware] Blocked tool '%s': arguments matched ATR rule %s.", + "[ATRValidationMiddleware] Blocked tool '%s': arguments matched ATR rule %s (%s, %s confidence).", context.function.name, - matched, + match.rule_id, + match.severity, + match.confidence, ) # Raise BEFORE call_next() so the tool is never executed. The matched rule id is # included for auditability. - raise MiddlewareTermination(f"ATR validation blocked tool '{context.function.name}' (rule: {matched})") + raise MiddlewareTermination( + f"ATR validation blocked tool '{context.function.name}' (rule: {match.rule_id})" + ) - logger.info("[ATRValidationMiddleware] Tool '%s' passed ATR validation.", context.function.name) + if match is not None: + # Medium or low severity: surface the finding for review, but let the tool run. + logger.warning( + "[ATRValidationMiddleware] Tool '%s' matched ATR rule %s (%s, %s confidence); not blocking.", + context.function.name, + match.rule_id, + match.severity, + match.confidence, + ) + else: + logger.info("[ATRValidationMiddleware] Tool '%s' passed ATR validation.", context.function.name) await call_next()