Skip to content

feat(guardrails): support during_call and MCP tool call modes in the Airia guardrail - #49709

Draft
airia-open-source wants to merge 6 commits into
BerriAI:mainfrom
airia-open-source:airia_guardrail_tool_hooks
Draft

airia-open-source wants to merge 6 commits into
BerriAI:mainfrom
airia-open-source:airia_guardrail_tool_hooks

Conversation

@airia-open-source

@airia-open-source airia-open-source commented Nov 2, 2026 •

Copy link
Copy Markdown
Contributor

The Airia guardrail now runs on two more points: while the model call is in flight (during_call), and around MCP tool calls the proxy makes (pre_mcp_call checks a tool's arguments before it runs, post_mcp_call checks its result). MCP evaluations tell the Airia gateway which tool they are for, so Airia applies the tenant's tool-call and tool-response filters to them instead of its prompt and response filters.

Details

TLDR

Problem this solves:

  • The Airia guardrail only supported pre_call and post_call
  • MCP tool arguments and results reached tools and users unchecked
  • No way to check a prompt in parallel with the model call

How it solves it:

  • Adds during_call, pre_mcp_call and post_mcp_call to the supported modes
  • during_call blocks on a block or a redact, since the call is already running
  • MCP evaluations send mcp_tool_name so Airia picks the tool filters
  • That marker comes from the hook event, never from the client's request body

User Flow

Before: a proxy admin who wants Airia to check MCP tool traffic cannot turn it on

  1. They add mode: [pre_call, post_call, pre_mcp_call, post_mcp_call] to the airia guardrail and start the proxy
  2. Startup rejects the config: pre_mcp_call is not a supported mode for airia
  3. A tool call through POST /mcp-rest/tools/call reaches the tool with the user's email address in its arguments

After: the same config checks MCP tool arguments and results

  1. They add the same modes and start the proxy
  2. The proxy starts with the guardrail active on all four points
  3. POST /mcp-rest/tools/call for send_email with an email in body: the tool receives [EmailAddress1]
  4. A tool whose result carries a phone number comes back to the caller redacted
  5. A tool call whose arguments carry an AWS key returns 400 and the tool never runs

Pre-Submission checklist

  • I have added meaningful tests
  • The handful of test files covering my change pass locally: pytest tests/unit/proxy/guardrails/guardrail_hooks/test_airia.py → 61 passed
  • My PR passes all required CI/CD checks
  • My PR's scope is as isolated as possible; it only solves 1 specific problem
  • Every automated review on my PR's latest commit has finished, and I have addressed or replied to each of its findings before requesting a maintainer review

Screenshots / Proof of Fix

Run against a live local Airia stack, no mocks: LiteLLM proxy → Airia AI Gateway (AiriaLLC/airia-ai-gateway#1317 at 6e1f225c) → Airia Platform and Policy Engine (origin/main, 2026-11-02) → the real PII and secrets detectors. The model is an echo server and the MCP server is a small tool server; both record what they actually received, so each case is checked on the far side of the proxy, not from the guardrail's own answer.

Shared setup:

  • Airia guardrail A51-317 LiteLLM demo on the tenant, targeted at LiteLLM: Standard PII (email address, phone number) → Redact & Replace; Secrets Detection → Block with the message Secret blocked by A51-317 demo policy. Inspection points stored on both filters: Input, Output, Tool call, Tool response.
  • Proxy config (After): guardrail: airia, mode: [pre_call, post_call, pre_mcp_call, post_mcp_call], default_on: true; the during_call cases use mode: [during_call, post_call].
  • MCP calls go through POST /mcp-rest/tools/call with {"server_id":"harness_tools","name":...,"arguments":{...}}; GET /_seen on the tool server lists what each tool received.

Airia Platform, the guardrail on the LiteLLM surface:

Airia Platform guardrails list, target LiteLLM
Airia Platform guardrail editor, LiteLLM surface
Airia Platform Standard PII filter, Redact & Replace

LiteLLM Dashboard, Test Guardrail panel (After; the airia guardrail lists all four modes):

LiteLLM Test Guardrail, email redacted
LiteLLM Test Guardrail, secret blocked

Before (f8f12a2)

Start the proxy with the new modes

  1. Start the proxy with mode: [during_call, post_call], then with mode: [pre_call, post_call, pre_mcp_call, post_mcp_call]
  2. Skipping guardrail 'airia': invalid configuration, proxy is starting WITHOUT this guardrail: Event hook GuardrailEventHooks.during_call is not in the supported event hooks [pre_call, post_call], and the same with pre_mcp_call for the second config

Chat: allowed, redacted, blocked

  1. Start the proxy with mode: [pre_call, post_call], the only modes supported here
  2. what is the capital of France? → 200, the model received it unchanged
  3. Email me at ada@example.com → 200, the model received Email me at [EmailAddress1]
  4. my aws key is AKIAIOSFODNN7EXAMPLE → 400 Secret blocked by A51-317 demo policy., the model was called 0 times

during_call: a secret in the prompt

  1. Config [during_call, post_call], POST /v1/chat/completions with my aws key is AKIAIOSFODNN7EXAMPLE
  2. 200 echo: my aws key is AKIAIOSFODNN7EXAMPLE: the guardrail was dropped at startup, nothing was checked

during_call: an email, which the gateway would redact

  1. Same config, POST /v1/chat/completions with Email me at ada@example.com
  2. 200 echo: Email me at ada@example.com

pre_mcp_call: an email in the arguments

  1. Config [pre_call, post_call, pre_mcp_call, post_mcp_call], call send_email with {"to":"ada@example.com","body":"hello"}
  2. 200; /_seen: {"tool":"send_email","to":"ada@example.com","body":"hello"}, 1 execution, nothing redacted

pre_mcp_call: an AWS key in the arguments

  1. Same config, call run_sql with {"query":"select 1 -- AKIAIOSFODNN7EXAMPLE"}
  2. 200; /_seen shows the query reached the tool, 1 execution

post_mcp_call: contact details in the result

  1. Same config, call lookup_contact with {"name":"Ada"}
  2. 200 Ada: reach ada@example.com or call 415-555-0199, nothing redacted

No Tool call or Tool response filters configured

  1. Same config, the guardrail's points set to Input and Output only, then the send_email, run_sql and lookup_contact calls above
  2. Same as the three cases above: the guardrail is not running, so everything reaches the tool and the caller untouched

After (0e1cb53)

Start the proxy with the new modes

  1. Start the proxy with mode: [during_call, post_call], then with mode: [pre_call, post_call, pre_mcp_call, post_mcp_call]
  2. Both start; GET /guardrails/list and the dashboard show airia with pre_call, post_call, pre_mcp_call, post_mcp_call

Chat: allowed, redacted, blocked

  1. Same prompts with the four-mode config
  2. what is the capital of France? → 200, the model received it unchanged
  3. Email me at ada@example.com → 200, the model received Email me at [EmailAddress1]
  4. my aws key is AKIAIOSFODNN7EXAMPLE → 400 Guardrail raised an exception, Guardrail: airia, Message: Secret blocked by A51-317 demo policy., the model was called 0 times

during_call: a secret in the prompt

  1. Config [during_call, post_call], POST /v1/chat/completions with my aws key is AKIAIOSFODNN7EXAMPLE
  2. 400 Guardrail raised an exception, Guardrail: airia, Message: Secret blocked by A51-317 demo policy.; the model was called once, in parallel with the check, and its reply was discarded

during_call: an email, which the gateway would redact

  1. Same config, POST /v1/chat/completions with Email me at ada@example.com
  2. 400 Guardrail raised an exception, Guardrail: airia, Message: Blocked by your organization's content policy.: the call is already running, so a redaction cannot be applied and comes back as a block

pre_mcp_call: an email in the arguments

  1. Four-mode config, call send_email with {"to":"ada@example.com","body":"hello"}
  2. 200; /_seen: {"tool":"send_email","to":"[EmailAddress1]","body":"hello"}, 1 execution

pre_mcp_call: an AWS key in the arguments

  1. Same config, call run_sql with {"query":"select 1 -- AKIAIOSFODNN7EXAMPLE"}
  2. 400 guardrail_violation ... Secret blocked by A51-317 demo policy.; /_seen is empty, the tool never ran

post_mcp_call: contact details in the result

  1. Same config, call lookup_contact with {"name":"Ada"}
  2. 200 Ada: reach [EmailAddress1] or call 415-555-0199; the email is redacted. The phone number is not: the detector does not flag the fictional 555 number, which is a detector matter, not this change

No Tool call or Tool response filters configured

  1. Same config, the guardrail's points set to Input and Output only (done through the Platform API), then the same three calls
  2. All three reach the tool and the caller untouched: the email arrives as ada@example.com, the AWS key query runs, the result is not redacted. MCP traffic is checked only under Tool call and Tool response filters, never under the prompt and response filters

Type

🆕 New Feature

Caveats (if any)

Medium

  • Needs the matching Airia gateway release
    • An older gateway ignores mcp_tool_name and checks MCP traffic with prompt/response filters
  • pre_mcp_call also runs when the proxy scans a server's tool catalog, which is how LiteLLM treats every guardrail in that mode (MCP_GUARDRAIL_CALL_TYPES includes list_mcp_tools; tool_catalog_guard runs on each tools/list)
    • Each listed tool costs one Airia evaluation, and its description and schema descriptions are sent as texts, ahead of any arguments
    • A Tool Call filter that matches text in a tool's description therefore hides that tool from tools/list and blocks calls to it whatever the arguments
    • While the Airia gateway is slow or unreachable, the scan fails closed and the listing comes back empty

Low

  • during_mcp_call is intentionally not supported

… a tool call

Cover the cases the marker must not appear in: during_call and post_call
evaluations, and client request data shaped like an MCP call (a top-level
mcp_tool_name, mcp_tool_call_metadata, or metadata.mcp_tool_name). Also pin
the fallback name used when an MCP evaluation carries none, and that a hook
that does not run leaves no marker for the next evaluation.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant