Uh oh!
There was an error while loading. Please reload this page.
feat(guardrails): Azure Content Safety Prompt Shield — P1 (#379) - #423
Conversation
…ent_safety) — P1
Adds a new guardrail kind that calls the Azure AI Content Safety
Prompt Shield API to detect jailbreak and indirect injection attacks.
Changes:
- `aisix-core`: new `AzureContentSafetyConfig` struct (`endpoint`,
`api_key`, `timeout_ms`) + `GuardrailKind::AzureContentSafety`
variant (serde tag `"azure_content_safety"`); exported from `models/mod.rs`.
- `aisix-guardrails`: new `prompt_shield.rs` implementing the `Guardrail`
trait via `reqwest`; POSTs to `/contentsafety/text:shieldPrompt?api-version=2024-09-01`;
auto-chunks prompts > 10 000 chars on whitespace boundaries;
maps 429 → `azure_cs_throttled`, 5xx/IO → `azure_cs_5xx`,
timeout → `azure_cs_timeout` with the same `fail_open` semantics as
the Bedrock kind.
- `aisix-guardrails/Cargo.toml`: `azure-content-safety` feature gate (default on)
pulling in `reqwest` from the workspace.
- `schemas/resources/guardrail.schema.json`: regenerated (new oneOf branch).
- 17 unit + wiremock tests cover the happy path, all failure modes,
hook-point gating, chunking, and the auth-header contract.
Wire shape (kine → DP):
{ "kind": "azure_content_safety",
"endpoint": "https://<resource>.cognitiveservices.azure.com",
"api_key": "<decrypted-plaintext>",
"timeout_ms": 5000 }
API reference: https://learn.microsoft.com/en-us/azure/ai-services/content-safety/referenceWarning Review limit reached
More reviews will be available in 11 minutes and 42 seconds. Learn how PR review limits work. Your organization has run out of usage credits. Purchase more in the billing tab. ⌛ How to resolve this issue?After more reviews become available, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans include higher PR review limits than trial, open-source, and free plans. In all cases, reviews become available again over time. During sustained high-volume PR review activity, CodeRabbit may temporarily slow when the next review becomes available. Please see our Fair Usage Limits Policy for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Free Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughThis PR adds an optional Azure AI Content Safety guardrail: new ChangesAzure Content Safety Guardrail
Sequence DiagramsequenceDiagram
participant Guardrail as check_input/output
participant PromptShield as PromptShieldGuardrail::shield
participant Chunker as chunk_text
participant AzureAPI as Azure shieldPrompt
participant FailureHandler as handle_failure
Guardrail->>PromptShield: collected text + hook_point
PromptShield->>Chunker: split into ~10k chunks
loop for each chunk
PromptShield->>AzureAPI: POST shieldPrompt (Ocp-Apim-Subscription-Key)
AzureAPI-->>PromptShield: response (attackDetected, analyses) or error/status
alt attackDetected == true
PromptShield-->>Guardrail: Block
else error/timeout
PromptShield->>FailureHandler: AcsFailure
FailureHandler-->>PromptShield: Bypass tag or Block (per fail_open)
end
end
PromptShield-->>Guardrail: Allow or Block or Bypass verdict
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Note 🎁 Summarized by CodeRabbit FreeYour organization has reached its limit of developer seats under the Pro Plan. For new users, CodeRabbit will generate a high-level summary and a walkthrough for each pull request. For a comprehensive line-by-line review, please add seats to your subscription by visiting https://app.coderabbit.ai/login.If you believe this is a mistake and have available seats, please assign one to the pull request author through the subscription management page using the link above. Comment |
MEDIUM-1: timeout_ms=0 doc contradicted actual behavior
- Fix doc comment in AzureContentSafetyConfig.timeout_ms: "0 = no timeout"
was wrong; Duration::ZERO fires on the first poll. Correct guidance:
use u32::MAX for an effectively unlimited timeout. Regenerate schema.
MEDIUM-2: happy-path test did not assert request body shape
- Add body_json + content-type matchers to
clean_input_returns_allow_and_sends_auth_header so a rename of
ShieldRequest.user_prompt or ShieldRequest.documents would catch a
real break rather than silently passing.
MEDIUM-3: 4xx (non-429) errors mislabeled as azure_cs_5xx
- Add AcsFailure::ConfigError + AcsFailure::ServerError, retiring the
ambiguous Other variant. 4xx gets bypass_tag "azure_cs_config_error";
5xx keeps "azure_cs_5xx". Log 4xx at error level (not warn) since
with fail_open=true a wrong api_key silently bypasses every request.
- Add two wiremock tests (HTTP 401 fail_open=true/false) to pin the
new ConfigError → azure_cs_config_error path.
- Extend bypass_tags_match_wire_contract to cover ServerError and
ConfigError variants.
- Update module-level behavior matrix table.
LOW-1: Fix inaccurate "disable pool timeout" comment in new()
LOW-2: chunk_text("") now returns [] not [""]; strengthen test assertion…2 (api-version pin) New MEDIUM from re-audit: ConfigError path fired tracing::error! in call_api() then tracing::warn! in handle_failure() for the same event, producing two log lines at different levels per request. Suppress the generic warn when the failure is ConfigError. LOW-2: pin api-version query parameter in the contract test so a version bump in SHIELD_PATH would be caught immediately. Add query_param matcher alongside the existing body_json and content-type assertions. LOW-1 (schema minimum=0) is intentional: the CP serializes timeout_ms with omitempty, so timeout_ms=0 is never forwarded to the DP — the DP defaults to 5000. The schema minimum=0 is correct for the CP→DP flow.
Summary
Implements the P1 guardrail kind: Azure AI Content Safety Prompt Shield (
kind=azure_content_safety).The Prompt Shield API detects two attack categories:
DP changes (
ai-gateway)aisix-core/src/models/guardrail.rsAzureContentSafetyConfigstruct (endpoint,api_key,timeout_ms); newGuardrailKind::AzureContentSafetyvariant (serde tag"azure_content_safety")aisix-core/src/models/mod.rsAzureContentSafetyConfigaisix-guardrails/src/prompt_shield.rsPromptShieldGuardrailimplementingGuardrail— POSTs to/contentsafety/text:shieldPrompt?api-version=2024-09-01; auto-chunks prompts > 10 000 chars on whitespace boundaries; maps 429 →azure_cs_throttled, 5xx/IO →azure_cs_5xx, timeout →azure_cs_timeoutwith the samefail_opensemantics as the Bedrock kindaisix-guardrails/src/build.rsAzureContentSafetyarm inbuild_one;azure-content-safetyfeature gate armaisix-guardrails/src/lib.rsPromptShieldGuardrail;#[cfg(feature = "azure-content-safety")] mod prompt_shieldaisix-guardrails/Cargo.tomlazure-content-safetyfeature gate (default on), pulls inreqwestfrom workspaceschemas/resources/guardrail.schema.jsonAzureContentSafetyConfigoneOfbranchWire shape (kine → DP)
{ "name": "my-shield", "kind": "azure_content_safety", "endpoint": "https://my-resource.cognitiveservices.azure.com", "api_key": "<decrypted-plaintext>", "timeout_ms": 5000, "fail_open": true }api_keyis decrypted by cp-api before kine write; the DP only ever holds plaintext in memory and never logs the key.API wire shape
Response:
{ "userPromptAnalysis": { "attackDetected": bool }, "documentsAnalysis": [] }Reference: https://learn.microsoft.com/en-us/azure/ai-services/content-safety/reference
Behavior matrix
fail_openattackDetected=false(all chunks)attackDetected=true(any chunk)azure_cs_timeoutazure_cs_throttledazure_cs_5xxChunking
Prompts > 10 000 chars are auto-split on whitespace boundaries. Each chunk ≤ 10 000 chars is sent as a separate API call. The first chunk returning
attackDetected=trueshort-circuits the rest.Test plan
prompt_shield::tests:chunk_textboundary conditions (exact limit, over limit, single oversized word, empty)handle_failureverdict paths (timeout/throttled + bothfail_openvalues)aisix-coremodel tests: parse round-trip,timeout_msdefaultDivergence from reference implementations
documents: []always sent (required field); we don't have the concept of "retrieved document grounding" at the gateway layer.tokio::time::timeout(same as ourlatency_mode=timedBedrock path) rather than reqwest's built-in timeout — keeps the failure mapping consistent with the rest of the codebase.Follow-ups
(cloud_safety_service, azure, prompt_shield)validation +marshalGuardrailKV+ API-key envelope encryptionazure_content_safetyguardrailsSummary by CodeRabbit
New Features
Tests
Documentation