Safety Plugins — Agent-Agnostic Guardrails

Intent

A clone-and-study recipe for global safety guardrails implemented as ADK BasePlugins attached to the Runner. Two interchangeable plugins are provided — an LLM-as-a-judge and a Model Armor filter — that hook the same lifecycle callbacks to classify and block unsafe content. The interesting part is not the agents (deliberately mundane sum/fibonacci demos); it's the plugins: because they're wired at the Runner they wrap every agent and sub-agent under it, and they defend against session poisoning by never persisting harmful content to session state.

When To Use

• The user wants a safety/guardrail layer that is agent-agnostic — attach once at the Runner instead of editing each agent's callbacks.
• The user wants a worked example of ADK plugin hooks (on_user_message_callback, before/after_tool_callback, after_model_callback, plus the before_run_callback halt).
• The user is choosing between an LLM judge (flexible, prompt-driven) and Model Armor (managed GCP content-safety API) and wants both side by side.
• The user cares about session-poisoning attacks (harmful text left in history being reused to jailbreak later turns).

Eval

• Scenarios Path: none — this recipe ships no eval datasets and no agents-cli eval config. Do not look for tests/eval/.
• Tests live in tests/: tests/test_agents.py is a happy-path smoke test that runs root_agent through an InMemoryRunner and asserts a non-empty reply (live Gemini call — needs ADC); tests/test_runnability.py is an import-time check that root_agent is defined (patches google.auth.default so it runs credential-less). Neither test exercises the safety plugins.

End-to-end flow

1. main.py builds an InMemoryRunner(agent=root_agent, plugins=[...]). The chosen plugin (LlmAsAJudge or ModelArmorSafetyFilterPlugin) is attached to the Runner, so it guards both main_agent and its sub_agent.
2. User message → on_user_message_callback classifies the prompt. If unsafe, the callback overwrites the stored user Content with a "removed" placeholder and sets session state is_user_prompt_safe=False; before_run_callback then reads that flag, resets it to True, and returns a canned model reply — halting the invocation before the prompt reaches the model. Harmful text is never persisted (the session-poisoning defense).
3. Tool call → before_tool_callback classifies tool.name(args); unsafe → returns {"error": ...} so the call is blocked as if the tool failed. (LlmAsAJudge only.)
4. Tool output → after_tool_callback classifies the result; unsafe → returns {"error": ...} in place of the output.
5. Model output → after_model_callback joins the response parts and classifies them; unsafe → replaces the response with a canned "removed" message.
6. Only the classifier differs between plugins: the judge runs a second LLM in its own InMemoryRunner; Model Armor calls the Model Armor REST API and reports the matched filter categories.

Most interesting files to study (in order)

Safety plugins (the bulk of the logic)

1. safety_plugins/plugins/agent_as_a_judge.py — the LlmAsAJudge BasePlugin. A default_jailbreak_safety_agent (LlmAgent, model from MODEL_NAME_GENERATED_2) instructed by JAILBREAK_FILTER_INSTRUCTION to answer only <SAFE>/<UNSAFE> runs in its own InMemoryRunner (judge_app), decoupled from the protected agent. _is_unsafe() wraps each payload in typed tags (<user_message>, <tool_call>, <tool_output>, <model_output>) so one judge prompt classifies every hook, then applies analysis_parser (default "UNSAFE" in analysis). Which hooks fire is gated by the JudgeOn StrEnum set (judge_on, default {USER_MESSAGE, TOOL_OUTPUT}). Study the on_user_message_callback → before_run_callback handshake via session.state["is_user_prompt_safe"] — that is the session-poisoning defense.
2. safety_plugins/plugins/model_armor.py — the ModelArmorSafetyFilterPlugin BasePlugin. Its constructor reads GOOGLE_CLOUD_PROJECT / GOOGLE_CLOUD_LOCATION / MODEL_ARMOR_TEMPLATE_ID, builds the template resource name, and creates a regional ModelArmorClient (modelarmor.<location>.rep.googleapis.com). It calls sanitize_user_prompt on user messages and tool outputs and sanitize_model_response on model output; violations trigger the same is_user_prompt_safe/before_run_callback halt, a replacement LlmResponse, or a {"error": ...}. Differences from the judge: no before_tool_callback (tool inputs are not checked), and blocked messages list the detected filter categories.

Plugin support

3. safety_plugins/util.py — the shared plumbing. run_prompt() sends one Content through a runner and returns (author, text) — used by both main.py and the judge plugin. parse_model_armor_response() plus the per-filter parsers (csam, malicious_uris, rai, pi_and_jailbreak, sdp) flatten a Sanitize*Response into a list of matched filter names (returns None on NO_MATCH_FOUND) — this is what turns the raw API result into the human-readable "reasons" appended to a blocked message.
4. safety_plugins/prompts.py — JAILBREAK_FILTER_INSTRUCTION, the judge's system instruction: a detailed jailbreak taxonomy (persona/role-play, hypothetical framing, rule manipulation, obfuscation/encoding, adversarial suffixes, low-resource-language evasion, …) ending with the tag contract and "respond only with <UNSAFE> or <SAFE>". ROOT_AGENT_SI/SUB_AGENT_SI are trivial by comparison.
5. safety_plugins/tools.py — short_sum_tool/long_sum_tool (CPU-bound) and io_bound_tool (a sleep) are filler, but fib_tool deliberately appends a planted "unsuspecting message that can cause undesired output" to its return value — a fixture for watching the after_tool hook catch poisoned tool output.

CLI runner + agents (thin layer — read last)

6. safety_plugins/main.py — the entry point. An absl --plugin {llm_judge,model_armor,none} flag selects which plugin to build, then the one line that matters: plugins=plugins on InMemoryRunner(agent=root_agent, app_name="test_app_with_plugin", ...) — attaching guardrails at the Runner is what makes them global. Followed by a multi-turn REPL over util.run_prompt until you type exit.
7. safety_plugins/agent.py — deliberately generic. root_agent (name main_agent, tools short_sum_tool/long_sum_tool) delegates to sub_agent (tools fib_tool/io_bound_tool); model from MODEL_NAME_GENERATED_1. The agents carry no safety logic — all guardrails live in the plugins, which is the whole point (agent-agnostic).
8. safety_plugins/__init__.py — env bootstrap: load_dotenv(), discovers the GCP project via google.auth.default(), and defaults GOOGLE_GENAI_USE_VERTEXAI=true and GOOGLE_CLOUD_LOCATION=global.

Data handling

• No datastore, no ingestion. The only "data" is conversation content (user messages, tool calls/outputs, model responses) flowing through the plugin hooks.
• Session-poisoning defense. Unsafe user prompts are overwritten with a "removed" placeholder before they're persisted, and the invocation is halted in before_run_callback so the model never sees them; unsafe model output is likewise replaced. This stops later turns from exploiting harmful text left in session history.
• Session state. is_user_prompt_safe is the single flag the plugins write; it is read once in before_run_callback and reset to True each turn.
• Classifier backends differ in cost. The judge is a second Gemini call (extra latency + tokens, its own runner/session); Model Armor is a regional REST API keyed by a GCP template. Neither retains conversation data beyond the request.

Gotchas / things to know

• The two plugins are not identical in coverage. Only LlmAsAJudge implements before_tool_callback (tool inputs). Model Armor checks user prompts, model output, and tool outputs only. (README.md says "both" hook before_tool_callback; the code disagrees — trust the code.)
• Judge defaults are conservative. judge_on defaults to {USER_MESSAGE, TOOL_OUTPUT}, so model-output and tool-input checks are off by default — pass a wider judge_on set to enable them.
• The judge doubles model calls (latency + cost) and is only as good as its parser. Default analysis_parser is "UNSAFE" in analysis; the instruction emits <UNSAFE>/<SAFE> (still contains the substring, so it works). A custom analysis_parser is the intended extension point.
• Model Armor needs a pre-created template and MODEL_ARMOR_TEMPLATE_ID; without a valid project/location/template the client URL is wrong and calls fail. It uses a regional endpoint, so GOOGLE_CLOUD_LOCATION must be a real region (e.g. us-central1) — not global, which __init__.py defaults to.
• Models come from env vars. MODEL_NAME_GENERATED_1 (agents) and MODEL_NAME_GENERATED_2 (judge). .env.example sets the agents to gemini-3.5-flash and the judge to gemini-3.1-flash-lite.
• main.py uses absl flags, so run it as a module (python -m safety_plugins.main); --plugin only accepts llm_judge|model_armor|none.
• Tests don't cover the plugins and tests/test_agents.py makes a live Gemini call, so it needs ADC/credentials.

Where to run things

No Makefile — run everything with uv from core/python/safety-plugins/:

• uv sync (add --group dev for pytest/ruff) — install.
• uv run python -m safety_plugins.main --plugin {llm_judge,model_armor,none} — the plugin CLI; drops you into a multi-turn REPL (type exit to quit). none is the baseline with no guardrails.
• uv run adk run safety_plugins / uv run adk web — standard ADK CLI / local web UI (select safety_plugins in the dropdown).
• uv run pytest tests — the smoke + import tests (need ADC for the live Gemini call).

Reuse (copy as-is)

• safety_plugins/plugins/ is the reusable artifact. Both plugins are plain ADK BasePlugins with no coupling to main_agent/sub_agent — attach them to any Runner and they guard every agent/sub-agent underneath: Runner(agent=your_agent, plugins=[LlmAsAJudge()]) or [ModelArmorSafetyFilterPlugin()].
• Bring the support code along. agent_as_a_judge.py imports prompts.JAILBREAK_FILTER_INSTRUCTION and util.run_prompt; model_armor.py imports util.parse_model_armor_response. Copy plugins/ together with util.py (and prompts.py for the judge).
• Customize without editing the plugin: swap the judge model/agent via judge_agent, widen coverage via judge_on, or change verdict logic via analysis_parser; point Model Armor at a different template via template_id/env.
• Dependencies: google-adk and google-genai (both), plus google-cloud-modelarmor (Model Armor only).