feat(guardrails): Azure Content Safety Prompt Shield — P1 (#379)

[moonming]

Summary

Implements the P1 guardrail kind: Azure AI Content Safety Prompt Shield (kind=azure_content_safety).

The Prompt Shield API detects two attack categories:

• Direct injection: jailbreak attempts in the user's own prompt
• Indirect injection: malicious instructions embedded in external documents

DP changes (ai-gateway)

File	Change
aisix-core/src/models/guardrail.rs	New AzureContentSafetyConfig struct (endpoint, api_key, timeout_ms); new GuardrailKind::AzureContentSafety variant (serde tag "azure_content_safety")
aisix-core/src/models/mod.rs	Export AzureContentSafetyConfig
aisix-guardrails/src/prompt_shield.rs	PromptShieldGuardrail implementing Guardrail — 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/src/build.rs	AzureContentSafety arm in build_one; azure-content-safety feature gate arm
aisix-guardrails/src/lib.rs	Export PromptShieldGuardrail; #[cfg(feature = "azure-content-safety")] mod prompt_shield
aisix-guardrails/Cargo.toml	azure-content-safety feature gate (default on), pulls in reqwest from workspace
schemas/resources/guardrail.schema.json	Regenerated — new AzureContentSafetyConfig oneOf branch

Wire 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_key is decrypted by cp-api before kine write; the DP only ever holds plaintext in memory and never logs the key.

API wire shape

POST {endpoint}/contentsafety/text:shieldPrompt?api-version=2024-09-01
Ocp-Apim-Subscription-Key: {api_key}

{ "userPrompt": "...", "documents": [] }

Response:

{ "userPromptAnalysis": { "attackDetected": bool }, "documentsAnalysis": [] }

Reference: https://learn.microsoft.com/en-us/azure/ai-services/content-safety/reference

Behavior matrix

API response	fail_open	Verdict
attackDetected=false (all chunks)	n/a	Allow
attackDetected=true (any chunk)	n/a	Block
timeout	true	Bypass azure_cs_timeout
timeout	false	Block
429 Throttling	true	Bypass azure_cs_throttled
429 Throttling	false	Block
5xx / IO error	true	Bypass azure_cs_5xx
5xx / IO error	false	Block

Chunking

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=true short-circuits the rest.

Test plan

• 17 unit + wiremock tests in prompt_shield::tests:
  • Bypass-tag contract (wire names must be stable)
  • chunk_text boundary conditions (exact limit, over limit, single oversized word, empty)
  • handle_failure verdict paths (timeout/throttled + both fail_open values)
  • Hook-point gating (output-only skips input; input-only skips output)
  • Wiremock: clean → Allow (+ auth header assertion), attack → Block, 5xx fail_open=true, 5xx fail_open=false, 429 throttled, timeout, long attack prompt, long clean prompt, output check, empty input
• 2 new aisix-core model tests: parse round-trip, timeout_ms default
• All 263 existing tests pass (178 aisix-core + 85 aisix-guardrails)
• E2E test against fakecloud (AISIX-Cloud PR — in progress)

Divergence from reference implementations

• No reference implementation (LiteLLM / Portkey) wraps Azure CS Prompt Shield today; wire shape derived directly from the Azure CS REST API docs (2024-09-01).
• documents: [] always sent (required field); we don't have the concept of "retrieved document grounding" at the gateway layer.
• Timeout handled via tokio::time::timeout (same as our latency_mode=timed Bedrock path) rather than reqwest's built-in timeout — keeps the failure mapping consistent with the rest of the codebase.

Follow-ups

• AISIX-Cloud CP PR: add (cloud_safety_service, azure, prompt_shield) validation + marshalGuardrailKV + API-key envelope encryption
• E2E test against fakecloud Azure CS stub
• Dashboard UI for creating/editing azure_content_safety guardrails

Summary by CodeRabbit

• New Features

  • Added Azure AI Content Safety guardrail for scanning chat inputs and outputs with chunking for long prompts.
  • New configuration options: endpoint, API key, and timeout (defaults to 5000ms); supports fail-open and fail-closed modes.

• Tests

  • Added unit and integration tests covering chunking, timeout behavior, success/failure mappings, and gating at hook points.

• Documentation

  • Schema updated to include the new guardrail kind and timeout default.

[coderabbitai]

Warning

Review limit reached

@moonming, we couldn't start this review because you've reached your PR review rate limit.

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 @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

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 configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Free

Run ID: 1b05770c-e246-47f3-a7bc-626207ab9710

📥 Commits

Reviewing files that changed from the base of the PR and between f17c713 and f52054a.

📒 Files selected for processing (1)

• crates/aisix-guardrails/src/prompt_shield.rs

📝 Walkthrough

Walkthrough

This PR adds an optional Azure AI Content Safety guardrail: new AzureContentSafetyConfig and GuardrailKind::AzureContentSafety, JSON schema support, crate feature/dependency for azure-content-safety (reqwest), builder wiring to construct PromptShieldGuardrail, and the PromptShield implementation with chunking, API calls, timeout/failure handling, and tests.

Changes

Azure Content Safety Guardrail

Layer / File(s)	Summary
Data model and schema contracts crates/aisix-core/src/models/guardrail.rs, crates/aisix-core/src/models/mod.rs, schemas/resources/guardrail.schema.json	AzureContentSafetyConfig with endpoint, api_key, and timeout_ms (defaults to 5000ms); GuardrailKind::AzureContentSafety variant added and serde rename_all = "snake_case" applied; re-export formatting adjusted; JSON Schema adds azure_content_safety variant; unit tests for explicit and default timeout deserialization.
Feature flag and dependency setup crates/aisix-guardrails/Cargo.toml, crates/aisix-guardrails/src/lib.rs	Adds optional reqwest dependency, defines azure-content-safety feature (enables dep:reqwest), includes it in default features, conditionally declares prompt_shield module and re-exports PromptShieldGuardrail behind the feature flag.
Builder integration crates/aisix-guardrails/src/build.rs	build_one matches GuardrailKind::AzureContentSafety and constructs PromptShieldGuardrail when feature enabled; when disabled returns BuildError::FeatureDisabled("azure-content-safety"); BuildError::FeatureDisabled made unconditionally available for multi-feature arms.
PromptShield implementation crates/aisix-guardrails/src/prompt_shield.rs	Adds PromptShieldGuardrail with HTTP client, endpoint normalization, stored API key, configurable timeout, and fail_open behavior; shield() chunks text and posts each chunk to Azure shieldPrompt, returning Block on any attackDetected=true; call_api() enforces timeout, maps 429/5xx/4xx/timeouts to AcsFailure; handle_failure() maps failures to Bypass (when fail_open) or Block (when not); includes serde request/response shapes, Guardrail trait impl gating by hook point, chunking helpers, and comprehensive tests (unit + tokio/wiremock integration).

Sequence Diagram

sequenceDiagram
  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

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

---

Note

🎁 Summarized by CodeRabbit Free

Your 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 @coderabbitai help to get the list of available commands and usage tips.}