diff --git a/crates/aisix-core/src/models/guardrail.rs b/crates/aisix-core/src/models/guardrail.rs index be317f85..0737757d 100644 --- a/crates/aisix-core/src/models/guardrail.rs +++ b/crates/aisix-core/src/models/guardrail.rs @@ -140,6 +140,122 @@ fn default_acs_timeout_ms() -> u32 { 5_000 } +/// Config block for `kind: "azure_content_safety_text_moderation"`. Calls +/// Azure AI Content Safety `text:analyze` for category-severity + blocklist +/// moderation on input and/or output (including streaming output). P2 +/// (PRD-09c §6 P2, #379). +/// +/// Reuses the P1 connection block (endpoint + api_key + timeout_ms). cp-api +/// projects only operator-set fields (omitempty), so every optional field +/// carries a serde default matching the cp-api validator's documented +/// default. Only `api_key` is a secret (decrypted by cp-api before kine +/// projection); every other field travels in the clear. +#[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct AzureContentSafetyTextModerationConfig { + /// Azure Cognitive Services resource endpoint. The DP appends + /// `/contentsafety/text:analyze?api-version=2024-09-01`. + pub endpoint: String, + /// Subscription key (`Ocp-Apim-Subscription-Key`). Plaintext in memory + /// only, never logged. + pub api_key: String, + /// HTTP call timeout (ms); `fail_open` / `output_fail_open` govern the + /// verdict when it elapses. See the `AzureContentSafetyConfig` note on + /// why `0` means "fire immediately", not "no timeout". + #[serde(default = "default_acs_timeout_ms")] + pub timeout_ms: u32, + + // --- moderation parameters --- + /// `FourSeverityLevels` (0,2,4,6; default) or `EightSeverityLevels` (0..7). + #[serde(default = "default_acs_output_type")] + pub output_type: String, + /// Categories to analyze. Defaults to all four. + #[serde(default = "default_acs_categories")] + pub categories: Vec, + /// General severity threshold; a category at or above it blocks. + #[serde(default = "default_acs_severity_threshold")] + pub severity_threshold: u8, + /// Per-category threshold overrides (take precedence over the general one). + #[serde(default)] + pub severity_threshold_by_category: std::collections::BTreeMap, + /// Azure CS blocklist names to match against. + #[serde(default)] + pub blocklist_names: Vec, + /// Forwarded to Azure's `haltOnBlocklistHit`. + #[serde(default)] + pub halt_on_blocklist_hit: bool, + /// Input-hook text selection: `concatenate_user_content` (default) or + /// `concatenate_all_content`. Ignored on the output hook. + #[serde(default = "default_acs_text_source")] + pub text_source: String, + + // --- streaming-output controls (consumed by aisix-proxy build_sse_stream) --- + /// `window` (sliding-window incremental release; default) or + /// `buffer_full` (whole-response hold-back). + #[serde(default = "default_acs_stream_processing_mode")] + pub stream_processing_mode: String, + /// Sliding-window size in chars (window mode); cp-api caps it at the + /// 10 000-char Azure limit. Default 10 000. + #[serde(default = "default_acs_window_size")] + pub window_size: u32, + /// Chars carried between windows so a span split across a boundary is + /// still caught. Default 256. + #[serde(default = "default_acs_window_overlap_size")] + pub window_overlap_size: u32, + /// Max bytes buffered in `buffer_full` mode before `on_buffer_exceeded` + /// applies. Default 262 144. + #[serde(default = "default_acs_max_buffer_bytes")] + pub max_buffer_bytes: u64, + /// `fail_closed` (default) or `fail_open` when the buffer cap is hit. + #[serde(default = "default_acs_on_buffer_exceeded")] + pub on_buffer_exceeded: String, + /// Fail-open policy for the OUTPUT hook. Defaults `false` (fail-closed) + /// so an Azure outage can't release unscanned model output. + #[serde(default)] + pub output_fail_open: bool, +} + +fn default_acs_output_type() -> String { + "FourSeverityLevels".to_owned() +} + +fn default_acs_categories() -> Vec { + vec![ + "Hate".to_owned(), + "Sexual".to_owned(), + "SelfHarm".to_owned(), + "Violence".to_owned(), + ] +} + +fn default_acs_severity_threshold() -> u8 { + 2 +} + +fn default_acs_text_source() -> String { + "concatenate_user_content".to_owned() +} + +fn default_acs_stream_processing_mode() -> String { + "window".to_owned() +} + +fn default_acs_window_size() -> u32 { + 10_000 +} + +fn default_acs_window_overlap_size() -> u32 { + 256 +} + +fn default_acs_max_buffer_bytes() -> u64 { + 262_144 +} + +fn default_acs_on_buffer_exceeded() -> String { + "fail_closed".to_owned() +} + /// Config block for `kind: "bedrock"`. Phase 1 stores the shape + /// passes it through `aisix-guardrails::build` which logs /// `bedrock not yet implemented` and skips the row. @@ -173,6 +289,11 @@ pub enum GuardrailKind { /// indirect injection attacks via the `/contentsafety/text:shieldPrompt` /// API. P1 (PRD-09c §6 P1). AzureContentSafety(AzureContentSafetyConfig), + /// Azure AI Content Safety Text Moderation. Category-severity + + /// blocklist moderation via the `/contentsafety/text:analyze` API, + /// on input and/or output (including streaming output). P2 + /// (PRD-09c §6 P2, #379). + AzureContentSafetyTextModeration(AzureContentSafetyTextModerationConfig), } /// Top-level `Guardrail` resource shape. Mirrors what cp-api writes @@ -554,6 +675,67 @@ mod tests { } } + #[test] + fn azure_text_moderation_kind_parses_with_defaults() { + // cp-api omits unset fields (omitempty); the DP must apply the + // documented defaults so a minimal row still moderates correctly. + let v = json!({ + "name": "moderate", + "kind": "azure_content_safety_text_moderation", + "endpoint": "https://my-resource.cognitiveservices.azure.com", + "api_key": "plaintext-key" + }); + let g: Guardrail = serde_json::from_value(v).unwrap(); + match g.config { + GuardrailKind::AzureContentSafetyTextModeration(ref c) => { + assert_eq!(c.timeout_ms, 5_000); + assert_eq!(c.output_type, "FourSeverityLevels"); + assert_eq!(c.severity_threshold, 2); + assert_eq!(c.categories.len(), 4); + assert_eq!(c.text_source, "concatenate_user_content"); + assert_eq!(c.stream_processing_mode, "window"); + assert_eq!(c.window_size, 10_000); + assert_eq!(c.window_overlap_size, 256); + assert_eq!(c.max_buffer_bytes, 262_144); + assert_eq!(c.on_buffer_exceeded, "fail_closed"); + assert!(!c.output_fail_open); + } + _ => panic!("expected AzureContentSafetyTextModeration variant"), + } + } + + #[test] + fn azure_text_moderation_kind_round_trips_set_fields() { + let v = json!({ + "name": "moderate", + "kind": "azure_content_safety_text_moderation", + "endpoint": "https://e.cognitiveservices.azure.com", + "api_key": "k", + "output_type": "EightSeverityLevels", + "categories": ["Hate", "Violence"], + "severity_threshold": 0, + "severity_threshold_by_category": { "Violence": 6 }, + "stream_processing_mode": "buffer_full", + "window_overlap_size": 0, + "output_fail_open": true + }); + let g: Guardrail = serde_json::from_value(v).unwrap(); + match g.config { + GuardrailKind::AzureContentSafetyTextModeration(ref c) => { + assert_eq!(c.output_type, "EightSeverityLevels"); + assert_eq!( + c.severity_threshold, 0, + "explicit 0 must survive (not defaulted to 2)" + ); + assert_eq!(c.severity_threshold_by_category.get("Violence"), Some(&6)); + assert_eq!(c.stream_processing_mode, "buffer_full"); + assert_eq!(c.window_overlap_size, 0, "explicit 0 overlap must survive"); + assert!(c.output_fail_open); + } + _ => panic!("expected AzureContentSafetyTextModeration variant"), + } + } + #[test] fn resource_trait_uses_name_and_guardrails_kind() { let mut g: Guardrail = serde_json::from_value(json!({ diff --git a/crates/aisix-core/src/models/mod.rs b/crates/aisix-core/src/models/mod.rs index cc883ddc..6ca7b166 100644 --- a/crates/aisix-core/src/models/mod.rs +++ b/crates/aisix-core/src/models/mod.rs @@ -30,9 +30,9 @@ pub mod snapshot; pub use apikey::ApiKey; pub use cache_policy::{AppliesTo, CacheBackend, CachePolicy}; pub use guardrail::{ - AzureContentSafetyConfig, BedrockAWSCredentials, BedrockConfig, BedrockLatencyMode, Guardrail, - GuardrailAttachment, GuardrailHookPoint, GuardrailKind, GuardrailScopeType, KeywordConfig, - KeywordPattern, + AzureContentSafetyConfig, AzureContentSafetyTextModerationConfig, BedrockAWSCredentials, + BedrockConfig, BedrockLatencyMode, Guardrail, GuardrailAttachment, GuardrailHookPoint, + GuardrailKind, GuardrailScopeType, KeywordConfig, KeywordPattern, }; pub use model::{ Adapter, BackgroundModelCheck, CooldownConfig, Model, DEFAULT_COOLDOWN_TRIGGER_STATUSES, diff --git a/crates/aisix-core/src/models/schema.rs b/crates/aisix-core/src/models/schema.rs index 1ef71629..39043918 100644 --- a/crates/aisix-core/src/models/schema.rs +++ b/crates/aisix-core/src/models/schema.rs @@ -409,7 +409,7 @@ fn guardrail_schema() -> Value { "enabled": { "type": "boolean" }, "hook_point": { "enum": ["input", "output", "both"] }, "fail_open": { "type": "boolean" }, - "kind": { "enum": ["keyword", "bedrock", "azure_content_safety"] } + "kind": { "enum": ["keyword", "bedrock", "azure_content_safety", "azure_content_safety_text_moderation"] } }, "oneOf": [ { @@ -451,6 +451,37 @@ fn guardrail_schema() -> Value { "api_key": { "type": "string", "minLength": 1 }, "timeout_ms": { "type": "integer", "minimum": 0, "maximum": 4_294_967_295u64 } } + }, + { + // kind=azure_content_safety_text_moderation — text:analyze + // category-severity + blocklist moderation. P2 (#379). + // Connection block matches azure_content_safety; the + // moderation + streaming params are optional (cp-api applies + // defaults + strict validation on write). + "type": "object", + "required": ["kind", "endpoint", "api_key"], + "properties": { + "kind": { "const": "azure_content_safety_text_moderation" }, + "endpoint": { "type": "string", "minLength": 1 }, + "api_key": { "type": "string", "minLength": 1 }, + "timeout_ms": { "type": "integer", "minimum": 0, "maximum": 4_294_967_295u64 }, + "output_type": { "enum": ["FourSeverityLevels", "EightSeverityLevels"] }, + "categories": { + "type": "array", + "items": { "enum": ["Hate", "Sexual", "SelfHarm", "Violence"] } + }, + "severity_threshold": { "type": "integer", "minimum": 0, "maximum": 7 }, + "severity_threshold_by_category": { "type": "object" }, + "blocklist_names": { "type": "array", "items": { "type": "string" } }, + "halt_on_blocklist_hit": { "type": "boolean" }, + "text_source": { "enum": ["concatenate_user_content", "concatenate_all_content"] }, + "stream_processing_mode": { "enum": ["window", "buffer_full"] }, + "window_size": { "type": "integer", "minimum": 1, "maximum": 10_000 }, + "window_overlap_size": { "type": "integer", "minimum": 0 }, + "max_buffer_bytes": { "type": "integer", "minimum": 1 }, + "on_buffer_exceeded": { "enum": ["fail_closed", "fail_open"] }, + "output_fail_open": { "type": "boolean" } + } } ], "$defs": { diff --git a/crates/aisix-guardrails/src/build.rs b/crates/aisix-guardrails/src/build.rs index 1b9f1435..603e497b 100644 --- a/crates/aisix-guardrails/src/build.rs +++ b/crates/aisix-guardrails/src/build.rs @@ -156,6 +156,23 @@ fn build_one( // Built without --features azure-content-safety. Skip + warn. Err(BuildError::FeatureDisabled("azure-content-safety")) } + #[cfg(feature = "azure-content-safety")] + GuardrailKind::AzureContentSafetyTextModeration(cfg) => { + // P2: HTTP-based text:analyze dispatcher. cp-api already + // decrypted the api_key at projection time; the config carries + // plaintext. Endpoint is per-row (each customer's own resource). + let g = crate::text_moderation::TextModerationGuardrail::new( + row.name.clone(), + cfg, + row.hook_point, + row.fail_open, + ); + Ok(Some(Arc::new(g))) + } + #[cfg(not(feature = "azure-content-safety"))] + GuardrailKind::AzureContentSafetyTextModeration(_) => { + Err(BuildError::FeatureDisabled("azure-content-safety")) + } } } diff --git a/crates/aisix-guardrails/src/chain.rs b/crates/aisix-guardrails/src/chain.rs index ed91f74e..aff9b03f 100644 --- a/crates/aisix-guardrails/src/chain.rs +++ b/crates/aisix-guardrails/src/chain.rs @@ -12,7 +12,7 @@ use std::sync::Arc; use aisix_gateway::{ChatFormat, ChatResponse}; use async_trait::async_trait; -use crate::{Guardrail, GuardrailVerdict}; +use crate::{Guardrail, GuardrailVerdict, StreamOutputPolicy}; #[derive(Clone)] pub struct GuardrailChain { @@ -56,6 +56,19 @@ impl Guardrail for GuardrailChain { self.guardrails.is_empty() } + /// The strictest streamed-output policy across the chain's members. + /// If any member wants hold-back, the whole stream holds back and + /// the full chain's `check_output` runs on the held content. + fn stream_output_policy(&self) -> StreamOutputPolicy { + self.guardrails + .iter() + .map(|g| g.stream_output_policy()) + .fold( + StreamOutputPolicy::EndOfStreamCheck, + StreamOutputPolicy::stricter, + ) + } + async fn check_input(&self, req: &ChatFormat) -> GuardrailVerdict { // `current` starts as a borrow; flips to Owned only if a Rewrite fires. // This keeps the common (no-rewrite) path allocation-free. diff --git a/crates/aisix-guardrails/src/lib.rs b/crates/aisix-guardrails/src/lib.rs index 0794142a..9805f978 100644 --- a/crates/aisix-guardrails/src/lib.rs +++ b/crates/aisix-guardrails/src/lib.rs @@ -29,6 +29,8 @@ mod keyword; mod length; #[cfg(feature = "azure-content-safety")] mod prompt_shield; +#[cfg(feature = "azure-content-safety")] +mod text_moderation; use aisix_gateway::{ChatFormat, ChatResponse}; use async_trait::async_trait; @@ -44,6 +46,8 @@ pub use keyword::{KeywordBlocklist, KeywordRule}; pub use length::MaxContentLength; #[cfg(feature = "azure-content-safety")] pub use prompt_shield::PromptShieldGuardrail; +#[cfg(feature = "azure-content-safety")] +pub use text_moderation::TextModerationGuardrail; /// What a guardrail decided about a request or response. /// @@ -113,6 +117,94 @@ impl GuardrailVerdict { } } +/// How a guardrail wants STREAMED output moderated. The proxy's SSE +/// builder queries [`Guardrail::stream_output_policy`] on the resolved +/// chain and applies the strictest member policy to decide whether to +/// hold streamed content back until it scans clean. +/// +/// `EndOfStreamCheck` is the pre-P2 behavior — chunks are forwarded +/// live and `check_output` runs once at end-of-stream (so a block frame +/// arrives *after* the content already reached the client). The +/// hold-back variants buffer content until it passes. +#[derive(Debug, Clone, PartialEq, Default)] +pub enum StreamOutputPolicy { + /// Forward live; check once at end-of-stream. No hold-back. Default. + #[default] + EndOfStreamCheck, + /// Sliding window: release a window of content only after it scans + /// clean; `overlap_chars` is carried between windows so a span split + /// across a boundary is still caught. + Window { + size_chars: usize, + overlap_chars: usize, + }, + /// Hold the whole response; scan once; release all or block. + /// `max_buffer_bytes` caps the hold; `on_exceeded_fail_open` decides + /// release-vs-block when the cap is exceeded. + BufferFull { + max_buffer_bytes: usize, + on_exceeded_fail_open: bool, + }, +} + +impl StreamOutputPolicy { + /// `true` when this policy holds streamed content back until it + /// scans clean (i.e. anything other than `EndOfStreamCheck`). + pub fn holds_back(&self) -> bool { + !matches!(self, StreamOutputPolicy::EndOfStreamCheck) + } + + /// Coarse strictness rank: more hold-back = higher. + fn rank(&self) -> u8 { + match self { + StreamOutputPolicy::EndOfStreamCheck => 0, + StreamOutputPolicy::Window { .. } => 1, + StreamOutputPolicy::BufferFull { .. } => 2, + } + } + + /// Pick the stricter of two policies (used to fold a chain into one). + /// Higher rank wins; ties break toward the tighter parameters + /// (smaller window, smaller buffer cap). + pub fn stricter(self, other: Self) -> Self { + use StreamOutputPolicy::*; + match self.rank().cmp(&other.rank()) { + std::cmp::Ordering::Less => other, + std::cmp::Ordering::Greater => self, + std::cmp::Ordering::Equal => match (self, other) { + ( + Window { + size_chars: a, + overlap_chars: oa, + }, + Window { + size_chars: b, + overlap_chars: ob, + }, + ) => Window { + size_chars: a.min(b), + overlap_chars: oa.max(ob), + }, + ( + BufferFull { + max_buffer_bytes: a, + on_exceeded_fail_open: fa, + }, + BufferFull { + max_buffer_bytes: b, + on_exceeded_fail_open: fb, + }, + ) => BufferFull { + max_buffer_bytes: a.min(b), + // fail-closed is stricter than fail-open. + on_exceeded_fail_open: fa && fb, + }, + (s, _) => s, + }, + } + } +} + /// Pluggable content-policy hook. Production wires `Arc` /// in `ProxyState`; tests construct in-memory chains directly. #[async_trait] @@ -138,6 +230,13 @@ pub trait Guardrail: Send + Sync + 'static { fn is_empty(&self) -> bool { false } + + /// How this guardrail wants streamed OUTPUT moderated. Default: + /// [`StreamOutputPolicy::EndOfStreamCheck`] (no hold-back, pre-P2 + /// behavior). Hold-back guardrails (Azure text moderation) override. + fn stream_output_policy(&self) -> StreamOutputPolicy { + StreamOutputPolicy::EndOfStreamCheck + } } #[cfg(test)] diff --git a/crates/aisix-guardrails/src/text_moderation.rs b/crates/aisix-guardrails/src/text_moderation.rs new file mode 100644 index 00000000..f2a7b9b5 --- /dev/null +++ b/crates/aisix-guardrails/src/text_moderation.rs @@ -0,0 +1,706 @@ +//! kind=azure_content_safety_text_moderation guardrail dispatcher — calls +//! Azure AI Content Safety `text:analyze` on chat input and/or output and +//! translates the category-severity + blocklist result into a +//! [`GuardrailVerdict`]. +//! +//! PRD-09c §6 P2 (#379). +//! +//! API reference (2024-09-01): +//! POST `{endpoint}/contentsafety/text:analyze?api-version=2024-09-01` +//! Source: +//! +//! Wire shape: +//! ```json +//! // Request +//! { "text": "...", "categories": ["Hate","Sexual","SelfHarm","Violence"], +//! "blocklistNames": [], "haltOnBlocklistHit": false, +//! "outputType": "FourSeverityLevels" } +//! // Response +//! { "categoriesAnalysis": [ { "category": "Hate", "severity": 2 }, ... ], +//! "blocklistsMatch": [ { "blocklistName": "...", "blocklistItemText": "..." } ] } +//! ``` +//! +//! Block decision: a category whose `severity` reaches its threshold +//! (per-category override → general threshold → default 2), OR a +//! non-empty `blocklistsMatch`. +//! +//! Streaming output is moderated separately in `aisix-proxy`'s +//! `build_sse_stream` using the `stream_processing_mode` / `window_*` +//! config; this dispatcher only implements the non-streaming +//! `check_input` / `check_output` hooks. +//! +//! NOTE: the HTTP transport (`chunk_text`, the `AcsFailure` buckets, the +//! fail-open verdict mapping, the `Ocp-Apim-Subscription-Key` + tokio +//! timeout call shape) is duplicated from `prompt_shield.rs`. Extracting a +//! shared `azure_common` module is tracked as a follow-up so this slice +//! stays surgical (it does not touch the shipped P1 dispatcher). + +use std::collections::BTreeMap; +use std::sync::Arc; +use std::time::Duration; + +use aisix_core::models::{AzureContentSafetyTextModerationConfig, GuardrailHookPoint}; +use aisix_gateway::{ChatFormat, ChatResponse, Role}; +use async_trait::async_trait; +use serde::{Deserialize, Serialize}; + +use crate::{Guardrail, GuardrailVerdict, StreamOutputPolicy}; + +/// Maximum characters per `text:analyze` call. Azure CS enforces a +/// 10 000-char limit on `text`. +const MAX_TEXT_CHARS: usize = 10_000; + +/// Path + query appended to the configured `endpoint`. +const ANALYZE_PATH: &str = "/contentsafety/text:analyze?api-version=2024-09-01"; + +/// One Azure Content Safety Text Moderation row, materialised into a +/// request-time dispatcher. +pub struct TextModerationGuardrail { + row_name: String, + endpoint: String, + api_key: String, + pub(crate) hook_point: GuardrailHookPoint, + /// Fail-open policy for the INPUT hook (from the outer `Guardrail`). + fail_open: bool, + /// Fail-open policy for the OUTPUT hook. Defaults to `false` + /// (fail-closed) so an Azure outage can't release unscanned model + /// output — otherwise output moderation is defeated by a timeout. + output_fail_open: bool, + pub(crate) timeout: Duration, + client: Arc, + + // --- moderation parameters --- + categories: Vec, + output_type: String, + severity_threshold: u8, + severity_threshold_by_category: BTreeMap, + blocklist_names: Vec, + halt_on_blocklist_hit: bool, + /// `concatenate_user_content` (default) scans only user messages on + /// the input hook; `concatenate_all_content` scans every message. + /// Ignored on the output hook (always the assistant message). + text_source: String, + + // --- streaming-output controls (surfaced via stream_output_policy; + // consumed by aisix-proxy's build_sse_stream) --- + stream_processing_mode: String, + window_size: u32, + window_overlap_size: u32, + max_buffer_bytes: u64, + on_buffer_exceeded: String, +} + +impl TextModerationGuardrail { + pub fn new( + row_name: impl Into, + cfg: &AzureContentSafetyTextModerationConfig, + hook_point: GuardrailHookPoint, + fail_open: bool, + ) -> Self { + let client = reqwest::Client::builder() + .build() + .expect("reqwest::Client::builder() failed; this should never happen"); + Self { + row_name: row_name.into(), + endpoint: cfg.endpoint.trim_end_matches('/').to_owned(), + api_key: cfg.api_key.clone(), + hook_point, + fail_open, + output_fail_open: cfg.output_fail_open, + timeout: Duration::from_millis(cfg.timeout_ms as u64), + client: Arc::new(client), + categories: cfg.categories.clone(), + output_type: cfg.output_type.clone(), + severity_threshold: cfg.severity_threshold, + severity_threshold_by_category: cfg.severity_threshold_by_category.clone(), + blocklist_names: cfg.blocklist_names.clone(), + halt_on_blocklist_hit: cfg.halt_on_blocklist_hit, + text_source: cfg.text_source.clone(), + stream_processing_mode: cfg.stream_processing_mode.clone(), + window_size: cfg.window_size, + window_overlap_size: cfg.window_overlap_size, + max_buffer_bytes: cfg.max_buffer_bytes, + on_buffer_exceeded: cfg.on_buffer_exceeded.clone(), + } + } + + /// Scan `text` in ≤10 000-char chunks. Returns `Block` on the first + /// chunk that crosses a category threshold or hits a blocklist; + /// `Allow` when every chunk is clean; the fail-open mapping on error. + async fn scan(&self, text: &str, fail_open: bool) -> GuardrailVerdict { + for chunk in chunk_text(text, MAX_TEXT_CHARS) { + match self.analyze(&chunk).await { + Ok(resp) => { + if let Some(reason) = self.violation_reason(&resp) { + return GuardrailVerdict::Block { reason }; + } + } + Err(failure) => return self.handle_failure(failure, fail_open), + } + } + GuardrailVerdict::Allow + } + + /// POST one chunk to `text:analyze` and return the parsed result. + async fn analyze(&self, text: &str) -> Result { + let url = format!("{}{}", self.endpoint, ANALYZE_PATH); + let body = AnalyzeRequest { + text, + categories: &self.categories, + blocklist_names: &self.blocklist_names, + halt_on_blocklist_hit: self.halt_on_blocklist_hit, + output_type: &self.output_type, + }; + + let future = self + .client + .post(&url) + .header("Ocp-Apim-Subscription-Key", &self.api_key) + .json(&body) + .send(); + + let resp = match tokio::time::timeout(self.timeout, future).await { + Err(_elapsed) => return Err(AcsFailure::Timeout), + Ok(Err(_e)) => return Err(AcsFailure::IoError), + Ok(Ok(r)) => r, + }; + + let status = resp.status(); + if status == reqwest::StatusCode::TOO_MANY_REQUESTS { + return Err(AcsFailure::Throttled); + } + if status.is_server_error() { + return Err(AcsFailure::ServerError); + } + if !status.is_success() { + tracing::error!( + row = %self.row_name, + http_status = status.as_u16(), + "azure content safety text:analyze returned 4xx — check endpoint and api_key configuration", + ); + return Err(AcsFailure::ConfigError); + } + + resp.json::() + .await + .map_err(|_| AcsFailure::ServerError) + } + + /// Apply the threshold + blocklist policy to one analyze response. + /// Returns the block reason, or `None` when the chunk is clean. + fn violation_reason(&self, resp: &AnalyzeResponse) -> Option { + for cat in &resp.categories_analysis { + let threshold = self + .severity_threshold_by_category + .get(&cat.category) + .copied() + .unwrap_or(self.severity_threshold); + if cat.severity >= threshold { + return Some(format!( + "azure content safety: {} severity {} >= threshold {} (row: {})", + cat.category, cat.severity, threshold, self.row_name + )); + } + } + if let Some(first) = resp.blocklists_match.first() { + return Some(format!( + "azure content safety: blocklist {:?} matched (row: {})", + first.blocklist_name, self.row_name + )); + } + None + } + + fn handle_failure(&self, failure: AcsFailure, fail_open: bool) -> GuardrailVerdict { + let tag = failure.bypass_tag(); + if !matches!(failure, AcsFailure::ConfigError) { + tracing::warn!( + row = %self.row_name, + failure = ?failure, + fail_open, + "azure content safety text moderation call failed", + ); + } + if fail_open { + GuardrailVerdict::Bypass { reason: tag.into() } + } else { + GuardrailVerdict::Block { + reason: format!("azure content safety unavailable ({tag})"), + } + } + } + + /// Collect the text the INPUT hook scans, honoring `text_source`. + fn collect_input_text(&self, req: &ChatFormat) -> String { + let all = self.text_source == "concatenate_all_content"; + req.messages + .iter() + .filter(|m| all || m.role == Role::User) + .map(|m| m.content.as_str()) + .filter(|s| !s.is_empty()) + .collect::>() + .join("\n") + } +} + +/// Failure cause buckets. `bypass_tag()` maps to the strings stored in +/// `usage_events.guardrail_bypassed_reason`; these match the P1 Prompt +/// Shield tags so operators filter both Azure kinds the same way. +#[derive(Debug)] +enum AcsFailure { + Timeout, + Throttled, + IoError, + ServerError, + ConfigError, +} + +impl AcsFailure { + fn bypass_tag(&self) -> &'static str { + match self { + Self::Timeout => "azure_cs_timeout", + Self::Throttled => "azure_cs_throttled", + Self::IoError | Self::ServerError => "azure_cs_5xx", + Self::ConfigError => "azure_cs_config_error", + } + } +} + +// --- serde shapes for the wire protocol ------------------------------------ + +#[derive(Serialize)] +struct AnalyzeRequest<'a> { + text: &'a str, + categories: &'a [String], + #[serde(rename = "blocklistNames")] + blocklist_names: &'a [String], + #[serde(rename = "haltOnBlocklistHit")] + halt_on_blocklist_hit: bool, + #[serde(rename = "outputType")] + output_type: &'a str, +} + +#[derive(Deserialize)] +struct AnalyzeResponse { + #[serde(rename = "categoriesAnalysis", default)] + categories_analysis: Vec, + #[serde(rename = "blocklistsMatch", default)] + blocklists_match: Vec, +} + +#[derive(Deserialize)] +struct CategoryAnalysis { + category: String, + severity: u8, +} + +#[derive(Deserialize)] +struct BlocklistMatch { + #[serde(rename = "blocklistName", default)] + blocklist_name: String, +} + +// --- Guardrail trait impl -------------------------------------------------- + +#[async_trait] +impl Guardrail for TextModerationGuardrail { + fn name(&self) -> &'static str { + "azure_content_safety_text_moderation" + } + + fn stream_output_policy(&self) -> StreamOutputPolicy { + match self.stream_processing_mode.as_str() { + "buffer_full" => StreamOutputPolicy::BufferFull { + max_buffer_bytes: self.max_buffer_bytes as usize, + on_exceeded_fail_open: self.on_buffer_exceeded == "fail_open", + }, + // "window" (default) and any unexpected value → sliding window. + _ => StreamOutputPolicy::Window { + size_chars: self.window_size as usize, + overlap_chars: self.window_overlap_size as usize, + }, + } + } + + async fn check_input(&self, req: &ChatFormat) -> GuardrailVerdict { + if !matches!( + self.hook_point, + GuardrailHookPoint::Input | GuardrailHookPoint::Both + ) { + return GuardrailVerdict::Allow; + } + let text = self.collect_input_text(req); + if text.is_empty() { + return GuardrailVerdict::Allow; + } + self.scan(&text, self.fail_open).await + } + + async fn check_output(&self, resp: &ChatResponse) -> GuardrailVerdict { + if !matches!( + self.hook_point, + GuardrailHookPoint::Output | GuardrailHookPoint::Both + ) { + return GuardrailVerdict::Allow; + } + let text = resp.message.content.clone(); + if text.is_empty() { + return GuardrailVerdict::Allow; + } + // Output uses its own fail policy (default fail-closed) so an + // Azure outage can't release unscanned model output. + self.scan(&text, self.output_fail_open).await + } +} + +/// Split `text` into chunks of at most `max_chars` characters on +/// whitespace boundaries. A single word over the limit is hard-truncated. +/// (Forked from `prompt_shield::chunk_text`; see the module note.) +fn chunk_text(text: &str, max_chars: usize) -> Vec { + if text.is_empty() { + return vec![]; + } + if text.chars().count() <= max_chars { + return vec![text.to_owned()]; + } + let mut chunks: Vec = Vec::new(); + let mut current = String::with_capacity(max_chars); + for word in text.split_whitespace() { + let word_chars = word.chars().count(); + let sep = if current.is_empty() { 0usize } else { 1 }; + if current.chars().count() + sep + word_chars > max_chars { + if !current.is_empty() { + chunks.push(std::mem::take(&mut current)); + } + if word_chars > max_chars { + chunks.push(word.chars().take(max_chars).collect()); + continue; + } + } + if !current.is_empty() { + current.push(' '); + } + current.push_str(word); + } + if !current.is_empty() { + chunks.push(current); + } + chunks +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use aisix_core::models::AzureContentSafetyTextModerationConfig; + use aisix_gateway::{ChatFormat, ChatMessage, ChatResponse, FinishReason, UsageStats}; + use serde_json::json; + use wiremock::matchers::{body_partial_json, header, method, path, query_param}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + use super::*; + + fn cfg(endpoint: &str) -> AzureContentSafetyTextModerationConfig { + // Mirrors what cp-api projects after applying its defaults. + serde_json::from_value(json!({ + "endpoint": endpoint, + "api_key": "test-key-abc", + "timeout_ms": 5_000, + })) + .unwrap() + } + + fn build(endpoint: &str, fail_open: bool) -> TextModerationGuardrail { + TextModerationGuardrail::new( + "wiremock-test", + &cfg(endpoint), + GuardrailHookPoint::Both, + fail_open, + ) + } + + fn req(msg: &str) -> ChatFormat { + ChatFormat::new("m", vec![ChatMessage::user(msg)]) + } + + fn resp(content: &str) -> ChatResponse { + ChatResponse { + id: "r".into(), + model: "m".into(), + message: ChatMessage::assistant(content), + finish_reason: FinishReason::Stop, + usage: UsageStats::new(0, 0), + } + } + + /// A clean analyze response (all categories below threshold). + fn clean_body() -> serde_json::Value { + json!({ + "categoriesAnalysis": [ + { "category": "Hate", "severity": 0 }, + { "category": "Violence", "severity": 0 } + ], + "blocklistsMatch": [] + }) + } + + // --- bypass-tag contract (shared with P1) --- + + #[test] + fn bypass_tags_match_wire_contract() { + assert_eq!(AcsFailure::Timeout.bypass_tag(), "azure_cs_timeout"); + assert_eq!(AcsFailure::Throttled.bypass_tag(), "azure_cs_throttled"); + assert_eq!(AcsFailure::IoError.bypass_tag(), "azure_cs_5xx"); + assert_eq!(AcsFailure::ServerError.bypass_tag(), "azure_cs_5xx"); + assert_eq!( + AcsFailure::ConfigError.bypass_tag(), + "azure_cs_config_error" + ); + } + + #[test] + fn chunk_text_exact_limit_is_not_split() { + let text: String = "a".repeat(MAX_TEXT_CHARS); + assert_eq!(chunk_text(&text, MAX_TEXT_CHARS).len(), 1); + } + + #[test] + fn stream_policy_reflects_config() { + // Defaults → sliding window 10000/256. + let g = build("http://unused", true); + assert_eq!( + g.stream_output_policy(), + StreamOutputPolicy::Window { + size_chars: 10_000, + overlap_chars: 256 + } + ); + // buffer_full mode surfaces the cap + on_exceeded policy. + let mut g2 = build("http://unused", true); + g2.stream_processing_mode = "buffer_full".to_owned(); + g2.max_buffer_bytes = 1000; + g2.on_buffer_exceeded = "fail_open".to_owned(); + assert_eq!( + g2.stream_output_policy(), + StreamOutputPolicy::BufferFull { + max_buffer_bytes: 1000, + on_exceeded_fail_open: true + } + ); + } + + // --- severity threshold logic (no HTTP) --- + + fn mk_resp(cats: &[(&str, u8)], blocklist: bool) -> AnalyzeResponse { + AnalyzeResponse { + categories_analysis: cats + .iter() + .map(|(c, s)| CategoryAnalysis { + category: (*c).to_owned(), + severity: *s, + }) + .collect(), + blocklists_match: if blocklist { + vec![BlocklistMatch { + blocklist_name: "corp-terms".to_owned(), + }] + } else { + vec![] + }, + } + } + + #[test] + fn default_threshold_blocks_at_two() { + let g = build("http://unused", true); + // severity 2 >= default 2 → block + assert!(g + .violation_reason(&mk_resp(&[("Hate", 2)], false)) + .is_some()); + // severity 0 < 2 → clean + assert!(g + .violation_reason(&mk_resp(&[("Hate", 0)], false)) + .is_none()); + } + + #[test] + fn per_category_override_is_independent() { + let mut g = build("http://unused", true); + g.severity_threshold = 2; + g.severity_threshold_by_category = BTreeMap::from([("Violence".to_owned(), 6u8)]); + // Violence severity 4 < its override 6 → clean, even though it + // would trip the general threshold of 2. + assert!(g + .violation_reason(&mk_resp(&[("Violence", 4)], false)) + .is_none()); + // Violence severity 6 >= override 6 → block. + assert!(g + .violation_reason(&mk_resp(&[("Violence", 6)], false)) + .is_some()); + // Hate has no override → general threshold 2 applies. + assert!(g + .violation_reason(&mk_resp(&[("Hate", 2)], false)) + .is_some()); + } + + #[test] + fn blocklist_match_blocks_regardless_of_severity() { + let g = build("http://unused", true); + assert!(g.violation_reason(&mk_resp(&[("Hate", 0)], true)).is_some()); + } + + // --- fail-open mapping (no HTTP) --- + + #[test] + fn timeout_fail_open_true_returns_bypass() { + let g = build("http://unused", true); + match g.handle_failure(AcsFailure::Timeout, true) { + GuardrailVerdict::Bypass { reason } => assert_eq!(reason, "azure_cs_timeout"), + other => panic!("expected Bypass, got {other:?}"), + } + } + + #[test] + fn output_defaults_fail_closed() { + // cp-api omits output_fail_open when unset → serde default false. + let g = build("http://unused", true); + assert!(!g.output_fail_open, "output must default to fail-closed"); + // An output-side failure with fail_open=false must Block. + assert!(g + .handle_failure(AcsFailure::Timeout, g.output_fail_open) + .is_block()); + } + + // --- text_source --- + + #[test] + fn user_content_source_skips_assistant_messages() { + let g = build("http://unused", true); + let mut c = ChatFormat::new( + "m", + vec![ + ChatMessage::user("user says hi"), + ChatMessage::assistant("assistant reply"), + ], + ); + c.messages.push(ChatMessage::user("more user text")); + let text = g.collect_input_text(&c); + assert!(text.contains("user says hi")); + assert!(text.contains("more user text")); + assert!( + !text.contains("assistant reply"), + "default concatenate_user_content must skip assistant messages" + ); + } + + #[test] + fn all_content_source_includes_assistant_messages() { + let mut g = build("http://unused", true); + g.text_source = "concatenate_all_content".to_owned(); + let c = ChatFormat::new( + "m", + vec![ + ChatMessage::user("user says hi"), + ChatMessage::assistant("assistant reply"), + ], + ); + let text = g.collect_input_text(&c); + assert!(text.contains("user says hi")); + assert!( + text.contains("assistant reply"), + "concatenate_all_content must include assistant messages" + ); + } + + // --- wiremock integration --- + + #[tokio::test] + async fn clean_input_returns_allow_and_sends_wire_shape() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/contentsafety/text:analyze")) + .and(query_param("api-version", "2024-09-01")) + .and(header("Ocp-Apim-Subscription-Key", "test-key-abc")) + .and(body_partial_json(json!({ + "text": "hello there", + "outputType": "FourSeverityLevels" + }))) + .respond_with(ResponseTemplate::new(200).set_body_json(clean_body())) + .expect(1) + .mount(&server) + .await; + + let g = build(&server.uri(), true); + assert_eq!( + g.check_input(&req("hello there")).await, + GuardrailVerdict::Allow + ); + } + + #[tokio::test] + async fn high_severity_input_returns_block() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/contentsafety/text:analyze")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "categoriesAnalysis": [ { "category": "Hate", "severity": 6 } ], + "blocklistsMatch": [] + }))) + .mount(&server) + .await; + + let g = build(&server.uri(), true); + assert!(g.check_input(&req("hateful content")).await.is_block()); + } + + #[tokio::test] + async fn http_5xx_fail_open_true_returns_bypass() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/contentsafety/text:analyze")) + .respond_with(ResponseTemplate::new(500)) + .mount(&server) + .await; + + let g = build(&server.uri(), true); + match g.check_input(&req("test")).await { + GuardrailVerdict::Bypass { reason } => assert_eq!(reason, "azure_cs_5xx"), + other => panic!("expected Bypass(azure_cs_5xx), got {other:?}"), + } + } + + #[tokio::test] + async fn output_5xx_fails_closed_by_default() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/contentsafety/text:analyze")) + .respond_with(ResponseTemplate::new(500)) + .mount(&server) + .await; + + // output_fail_open defaults false → an output-side 5xx must Block. + let g = build(&server.uri(), true); + assert!( + g.check_output(&resp("some model output")).await.is_block(), + "output hook must fail closed on Azure error by default" + ); + } + + #[tokio::test] + async fn high_severity_output_returns_block() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/contentsafety/text:analyze")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "categoriesAnalysis": [ { "category": "Violence", "severity": 4 } ], + "blocklistsMatch": [] + }))) + .mount(&server) + .await; + + let g = build(&server.uri(), true); + assert!(g.check_output(&resp("violent output")).await.is_block()); + } +} diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index 2d25e1b7..2cf0d4c3 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -2093,6 +2093,24 @@ where } else { None }; + // P2 (#379): per-guardrail streamed-output policy. EndOfStreamCheck + // (default / no guardrail) leaves the live-forward path below + // byte-for-byte unchanged. Window / BufferFull hold content back + // until it scans clean. + let stream_policy = output_guardrail + .as_ref() + .map(|ctx| ctx.chain.stream_output_policy()) + .unwrap_or_default(); + let hold_back = stream_policy.holds_back(); + // Rendered content events withheld from the wire until their + // window (or the whole response) scans clean. Hold-back path only. + let mut pending: Vec = Vec::new(); + // Content accumulated since the last window flush (Window mode); + // bounded to ~window_size, unlike content_buffer (whole response). + let mut window_buf = String::new(); + // Set when a BufferFull cap is exceeded with fail-open: stop + // holding and forward the remainder live. + let mut cap_released = false; // Accumulate the upstream's `usage` block + per-chunk metadata // across the stream. Providers typically populate `usage` on // the terminal chunk only; using "max" rather than "last" makes @@ -2110,7 +2128,7 @@ where let mut errored = false; let mut first_chunk_seen = false; while let Some(item) = upstream.next().await { - let ev = match item { + let (ev, is_error_ev) = match item { Ok(chunk) => { // Record TTFT on the first chunk carrying generated // output (content or tool calls). Skip role-only @@ -2132,15 +2150,19 @@ where if let Some(fr) = chunk.finish_reason.as_ref() { comp.finish_reason = finish_reason_label(fr); } - // Per #204: accumulate the assistant's content - // when an output guardrail is configured. Skip - // entirely when none is configured to avoid the - // allocation on the hot path. - if let (Some(buf), Some(text)) = ( - content_buffer.as_mut(), - chunk.delta.content.as_deref(), - ) { - buf.push_str(text); + // Accumulate the assistant's content for the output + // guardrail. Window mode keeps only the current window + // (bounded memory); other modes accumulate the whole + // response in content_buffer. No-op without a guardrail. + if let Some(text) = chunk.delta.content.as_deref() { + if matches!( + stream_policy, + aisix_guardrails::StreamOutputPolicy::Window { .. } + ) { + window_buf.push_str(text); + } else if let Some(buf) = content_buffer.as_mut() { + buf.push_str(text); + } } if let Some(u) = chunk.usage.as_ref() { if u.prompt_tokens > comp.prompt_tokens { @@ -2168,24 +2190,143 @@ where } let rendered = render_chunk(created, chunk, &client_facing_model); match serde_json::to_string(&rendered) { - Ok(json) => Event::default().data(json), + Ok(json) => (Event::default().data(json), false), Err(err) => { errored = true; - Event::default() - .event("error") - .data(error_frame_payload("internal_error", &err.to_string())) + ( + Event::default() + .event("error") + .data(error_frame_payload("internal_error", &err.to_string())), + true, + ) } } } Err(err) => { errored = true; let etype = err.error_type(); - Event::default() - .event("error") - .data(error_frame_payload(etype, &err.to_string())) + ( + Event::default() + .event("error") + .data(error_frame_payload(etype, &err.to_string())), + true, + ) } }; - yield Ok::<_, Infallible>(ev); + if is_error_ev || !hold_back || cap_released { + // Error frames, the EndOfStreamCheck path, and a released + // BufferFull cap all forward straight to the wire. On the + // hold-back path an error drops the held (unscanned) + // content via the `errored` skip at end-of-stream (fail + // closed). + yield Ok::<_, Infallible>(ev); + } else { + // Hold-back: withhold this content event until its window + // (or the whole response) scans clean. + pending.push(ev); + match &stream_policy { + aisix_guardrails::StreamOutputPolicy::Window { + size_chars, + overlap_chars, + } => { + if window_buf.chars().count() >= *size_chars { + if let Some(ctx) = output_guardrail.as_ref() { + let synthesized = { + let comp = guard.comp(); + aisix_gateway::ChatResponse { + id: comp.provider_request_id.clone(), + model: comp.provider_model_version.clone(), + message: aisix_gateway::ChatMessage::assistant( + window_buf.clone(), + ), + finish_reason: aisix_gateway::FinishReason::Stop, + usage: aisix_gateway::UsageStats::new( + comp.prompt_tokens, + comp.completion_tokens, + ), + } + }; + match ctx.chain.check_output(&synthesized).await { + aisix_guardrails::GuardrailVerdict::Block { reason } => { + tracing::warn!( + guardrail_hook = "output", + model = %ctx.model_name, + reason = %reason, + "guardrail blocked streaming response (window)", + ); + errored = true; + guard.comp().guardrail_blocked = true; + yield Ok::<_, Infallible>( + Event::default().event("error").data( + error_frame_payload( + "content_filter", + "response blocked by content policy", + ), + ), + ); + break; + } + aisix_guardrails::GuardrailVerdict::Bypass { reason } => { + let comp = guard.comp(); + if comp.bypass_reason.is_empty() { + comp.bypass_reason = reason; + } + } + _ => {} + } + // Clean (Allow / Bypass / Rewrite): release + // this window's events, then keep the + // trailing overlap as scan context for the + // next window (its events were already sent). + for e in pending.drain(..) { + yield Ok::<_, Infallible>(e); + } + // Clamp the retained overlap to cc-1 so a + // misconfigured overlap >= window can't keep + // the whole buffer and re-scan every + // subsequent token (cost/latency guard). + let cc = window_buf.chars().count(); + let keep = (*overlap_chars).min(cc.saturating_sub(1)); + window_buf = if keep > 0 { + window_buf.chars().skip(cc - keep).collect() + } else { + String::new() + }; + } + } + } + aisix_guardrails::StreamOutputPolicy::BufferFull { + max_buffer_bytes, + on_exceeded_fail_open, + } => { + let buffered = content_buffer.as_ref().map_or(0, |b| b.len()); + if buffered > *max_buffer_bytes { + if *on_exceeded_fail_open { + cap_released = true; + for e in pending.drain(..) { + yield Ok::<_, Infallible>(e); + } + } else { + tracing::warn!( + guardrail_hook = "output", + max_buffer_bytes = *max_buffer_bytes, + "streaming response exceeded max_buffer_bytes; failing closed", + ); + errored = true; + guard.comp().guardrail_blocked = true; + yield Ok::<_, Infallible>( + Event::default().event("error").data(error_frame_payload( + "content_filter", + "response blocked by content policy", + )), + ); + break; + } + } + } + aisix_guardrails::StreamOutputPolicy::EndOfStreamCheck => {} + } + } // Delivery is counted by the outer DeliveryCounter wrapper // at poll_next time, not here — async_stream's yield // suspends BEFORE the consumer has actually pulled, so @@ -2210,8 +2351,69 @@ where // signaled the failure via SSE error event. Running the // guardrail on the partial would only add a second error // frame, not retroactively redact the leak. - if !errored { - if let (Some(content), Some(ctx)) = (content_buffer.as_ref(), output_guardrail.as_ref()) { + if !errored && !cap_released { + if hold_back { + // Final scan of the held content (the last partial window, + // or — for BufferFull — the whole response), then release + // the held events if it scans clean. + if let Some(ctx) = output_guardrail.as_ref() { + let final_text = match &stream_policy { + aisix_guardrails::StreamOutputPolicy::Window { .. } => window_buf.clone(), + _ => content_buffer.clone().unwrap_or_default(), + }; + let blocked = if final_text.is_empty() { + false + } else { + let synthesized = { + let comp = guard.comp(); + aisix_gateway::ChatResponse { + id: comp.provider_request_id.clone(), + model: comp.provider_model_version.clone(), + message: aisix_gateway::ChatMessage::assistant(final_text), + finish_reason: aisix_gateway::FinishReason::Stop, + usage: aisix_gateway::UsageStats::new( + comp.prompt_tokens, + comp.completion_tokens, + ), + } + }; + match ctx.chain.check_output(&synthesized).await { + aisix_guardrails::GuardrailVerdict::Block { reason } => { + tracing::warn!( + guardrail_hook = "output", + model = %ctx.model_name, + reason = %reason, + "guardrail blocked streaming response", + ); + errored = true; + guard.comp().guardrail_blocked = true; + yield Ok::<_, Infallible>( + Event::default().event("error").data(error_frame_payload( + "content_filter", + "response blocked by content policy", + )), + ); + true + } + aisix_guardrails::GuardrailVerdict::Bypass { reason } => { + let comp = guard.comp(); + if comp.bypass_reason.is_empty() { + comp.bypass_reason = reason; + } + false + } + _ => false, + } + }; + if !blocked { + for e in pending.drain(..) { + yield Ok::<_, Infallible>(e); + } + } + } + } else if let (Some(content), Some(ctx)) = + (content_buffer.as_ref(), output_guardrail.as_ref()) + { let synthesized = aisix_gateway::ChatResponse { id: guard.comp().provider_request_id.clone(), model: guard.comp().provider_model_version.clone(), diff --git a/crates/aisix-proxy/src/lib.rs b/crates/aisix-proxy/src/lib.rs index 4ac2e204..c471adce 100644 --- a/crates/aisix-proxy/src/lib.rs +++ b/crates/aisix-proxy/src/lib.rs @@ -1915,6 +1915,326 @@ data: [DONE]\n\n"; ); } + /// P2 (#379): unlike the keyword guardrail above (EndOfStreamCheck, + /// which leaks pre-emitted chunks), `azure_content_safety_text_moderation` + /// uses a hold-back streaming policy — a blocked streaming response + /// NEVER puts the offending content on the wire. + #[tokio::test] + async fn streaming_text_moderation_blocks_and_holds_content_back_no_leak() { + let upstream = MockServer::start().await; + let sse = "\ +data: {\"id\":\"up-1\",\"model\":\"gpt-4o\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\"},\"finish_reason\":null}]}\n\n\ +data: {\"id\":\"up-1\",\"model\":\"gpt-4o\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"this is harmful text\"},\"finish_reason\":null}]}\n\n\ +data: {\"id\":\"up-1\",\"model\":\"gpt-4o\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n\ +data: [DONE]\n\n"; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(sse), + ) + .mount(&upstream) + .await; + + // Azure Content Safety mock returns high severity → block. + let acs = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/contentsafety/text:analyze")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "categoriesAnalysis": [{"category": "Hate", "severity": 6}], + "blocklistsMatch": [] + }))) + .mount(&acs) + .await; + + let hub = Arc::new(Hub::new()); + hub.register_specialized("openai", Arc::new(OpenAiBridge::new())); + let snap = seed_snapshot("my-gpt4", &["my-gpt4"], &upstream.uri()); + let state = build_state(snap, hub); + seed_guardrail( + &state.snapshot, + "g-textmod-output", + &format!( + r#"{{"name":"textmod","kind":"azure_content_safety_text_moderation","hook_point":"output","endpoint":"{}","api_key":"k"}}"#, + acs.uri() + ), + ); + let app = build_router(state); + + let body = serde_json::json!({"model":"my-gpt4","messages":[{"role":"user","content":"hi"}],"stream":true}); + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", "Bearer sk-caller") + .header("content-type", "application/json") + .body(Body::from(body.to_string())) + .unwrap(); + let resp = run(app, req).await; + assert_eq!(resp.status(), StatusCode::OK); + + let mut body_stream = resp.into_body().into_data_stream(); + let mut wire = Vec::new(); + while let Some(chunk) = body_stream.next().await { + wire.extend_from_slice(chunk.unwrap().as_ref()); + } + let wire_str = String::from_utf8(wire).expect("SSE bytes are utf8"); + + // The hold-back guarantee: the harmful content NEVER reached the + // wire (held in `pending`, dropped on block). + assert!( + !wire_str.contains("harmful text"), + "hold-back must keep blocked content off the wire; got:\n{wire_str}" + ); + assert!( + wire_str.contains("event: error"), + "blocked stream must emit `event: error`; got:\n{wire_str}" + ); + assert!( + !wire_str.contains("data: [DONE]"), + "blocked stream must omit [DONE]; got:\n{wire_str}" + ); + let idx = wire_str.find("event: error\n").unwrap(); + let data_line = wire_str[idx..] + .lines() + .find(|l| l.starts_with("data: ")) + .unwrap(); + let parsed: serde_json::Value = serde_json::from_str(&data_line["data: ".len()..]).unwrap(); + assert_eq!(parsed["error"]["type"], "content_filter"); + } + + /// Clean content is held back, scanned, then released in full + [DONE]. + #[tokio::test] + async fn streaming_text_moderation_releases_clean_content() { + let upstream = MockServer::start().await; + let sse = "\ +data: {\"id\":\"up-1\",\"model\":\"gpt-4o\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"perfectly fine answer\"},\"finish_reason\":null}]}\n\n\ +data: {\"id\":\"up-1\",\"model\":\"gpt-4o\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n\ +data: [DONE]\n\n"; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(sse), + ) + .mount(&upstream) + .await; + let acs = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/contentsafety/text:analyze")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "categoriesAnalysis": [{"category": "Hate", "severity": 0}], + "blocklistsMatch": [] + }))) + .mount(&acs) + .await; + let hub = Arc::new(Hub::new()); + hub.register_specialized("openai", Arc::new(OpenAiBridge::new())); + let snap = seed_snapshot("my-gpt4", &["my-gpt4"], &upstream.uri()); + let state = build_state(snap, hub); + seed_guardrail( + &state.snapshot, + "g-textmod-clean", + &format!( + r#"{{"name":"textmod","kind":"azure_content_safety_text_moderation","hook_point":"output","endpoint":"{}","api_key":"k"}}"#, + acs.uri() + ), + ); + let app = build_router(state); + let body = serde_json::json!({"model":"my-gpt4","messages":[{"role":"user","content":"hi"}],"stream":true}); + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", "Bearer sk-caller") + .header("content-type", "application/json") + .body(Body::from(body.to_string())) + .unwrap(); + let resp = run(app, req).await; + assert_eq!(resp.status(), StatusCode::OK); + let mut body_stream = resp.into_body().into_data_stream(); + let mut wire = Vec::new(); + while let Some(c) = body_stream.next().await { + wire.extend_from_slice(c.unwrap().as_ref()); + } + let wire_str = String::from_utf8(wire).expect("utf8"); + assert!( + wire_str.contains("perfectly fine answer"), + "clean content must be released after the scan; got:\n{wire_str}" + ); + assert!( + wire_str.contains("data: [DONE]"), + "clean stream must end with [DONE]; got:\n{wire_str}" + ); + assert!( + !wire_str.contains("event: error"), + "clean stream must not emit an error frame; got:\n{wire_str}" + ); + } + + /// Drive a streaming chat through a seeded text-moderation guardrail + /// and return the raw SSE wire bytes. `guardrail_cfg` is the full + /// guardrail JSON (with the ACS mock endpoint already substituted). + async fn run_textmod_stream(guardrail_cfg: &str, upstream_sse: &str) -> String { + let upstream = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(upstream_sse.to_owned()), + ) + .mount(&upstream) + .await; + let hub = Arc::new(Hub::new()); + hub.register_specialized("openai", Arc::new(OpenAiBridge::new())); + let snap = seed_snapshot("my-gpt4", &["my-gpt4"], &upstream.uri()); + let state = build_state(snap, hub); + seed_guardrail(&state.snapshot, "g-tm", guardrail_cfg); + let app = build_router(state); + let body = serde_json::json!({"model":"my-gpt4","messages":[{"role":"user","content":"hi"}],"stream":true}); + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", "Bearer sk-caller") + .header("content-type", "application/json") + .body(Body::from(body.to_string())) + .unwrap(); + let resp = run(app, req).await; + assert_eq!(resp.status(), StatusCode::OK); + let mut body_stream = resp.into_body().into_data_stream(); + let mut wire = Vec::new(); + while let Some(c) = body_stream.next().await { + wire.extend_from_slice(c.unwrap().as_ref()); + } + String::from_utf8(wire).expect("utf8") + } + + async fn acs_mock(severity: u8) -> MockServer { + let acs = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/contentsafety/text:analyze")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "categoriesAnalysis": [{"category": "Hate", "severity": severity}], + "blocklistsMatch": [] + }))) + .mount(&acs) + .await; + acs + } + + fn two_content_chunks(a: &str, b: &str) -> String { + format!( + "data: {{\"id\":\"u\",\"model\":\"gpt-4o\",\"choices\":[{{\"index\":0,\"delta\":{{\"content\":\"{a}\"}},\"finish_reason\":null}}]}}\n\n\ +data: {{\"id\":\"u\",\"model\":\"gpt-4o\",\"choices\":[{{\"index\":0,\"delta\":{{\"content\":\"{b}\"}},\"finish_reason\":null}}]}}\n\n\ +data: {{\"id\":\"u\",\"model\":\"gpt-4o\",\"choices\":[{{\"index\":0,\"delta\":{{}},\"finish_reason\":\"stop\"}}]}}\n\n\ +data: [DONE]\n\n" + ) + } + + /// H1: Window mode blocks MID-STREAM (small window so the first window + /// trips before end-of-stream) and leaks nothing. + #[tokio::test] + async fn streaming_text_moderation_window_blocks_mid_stream() { + let acs = acs_mock(6).await; + let cfg = format!( + r#"{{"name":"tm","kind":"azure_content_safety_text_moderation","hook_point":"output","endpoint":"{}","api_key":"k","stream_processing_mode":"window","window_size":5,"window_overlap_size":1}}"#, + acs.uri() + ); + let wire = run_textmod_stream(&cfg, &two_content_chunks("hello ", "world!")).await; + assert!( + !wire.contains("hello"), + "mid-stream block must not leak window content; got:\n{wire}" + ); + assert!( + !wire.contains("world"), + "mid-stream block must not leak later content; got:\n{wire}" + ); + assert!( + wire.contains("event: error"), + "expected content_filter frame; got:\n{wire}" + ); + assert!( + !wire.contains("data: [DONE]"), + "blocked stream omits [DONE]; got:\n{wire}" + ); + } + + /// H1: Window mode releases multiple clean windows (exercises the + /// mid-stream flush + overlap retention), ending with [DONE]. + #[tokio::test] + async fn streaming_text_moderation_window_releases_clean_multiwindow() { + let acs = acs_mock(0).await; + let cfg = format!( + r#"{{"name":"tm","kind":"azure_content_safety_text_moderation","hook_point":"output","endpoint":"{}","api_key":"k","stream_processing_mode":"window","window_size":5,"window_overlap_size":2}}"#, + acs.uri() + ); + let wire = run_textmod_stream(&cfg, &two_content_chunks("hello ", "world!")).await; + assert!( + wire.contains("hello"), + "clean windows must be released; got:\n{wire}" + ); + assert!( + wire.contains("world"), + "all clean content must be released; got:\n{wire}" + ); + assert!( + wire.contains("data: [DONE]"), + "clean stream ends with [DONE]; got:\n{wire}" + ); + assert!( + !wire.contains("event: error"), + "clean stream emits no error; got:\n{wire}" + ); + } + + /// H1: BufferFull cap exceeded with fail_closed → block, no leak. + #[tokio::test] + async fn streaming_text_moderation_buffer_full_cap_fail_closed_blocks() { + let acs = acs_mock(0).await; // severity irrelevant — the cap trips first + let cfg = format!( + r#"{{"name":"tm","kind":"azure_content_safety_text_moderation","hook_point":"output","endpoint":"{}","api_key":"k","stream_processing_mode":"buffer_full","max_buffer_bytes":4,"on_buffer_exceeded":"fail_closed"}}"#, + acs.uri() + ); + let wire = run_textmod_stream(&cfg, &two_content_chunks("abcd", "efghij")).await; + assert!( + !wire.contains("abcd"), + "fail-closed cap must not leak buffered content; got:\n{wire}" + ); + assert!( + wire.contains("event: error"), + "cap fail-closed must emit content_filter; got:\n{wire}" + ); + assert!( + !wire.contains("data: [DONE]"), + "cap-blocked stream omits [DONE]; got:\n{wire}" + ); + } + + /// H1: BufferFull cap exceeded with fail_open → release held + forward + /// the rest live, ending with [DONE]. + #[tokio::test] + async fn streaming_text_moderation_buffer_full_cap_fail_open_releases() { + let acs = acs_mock(0).await; + let cfg = format!( + r#"{{"name":"tm","kind":"azure_content_safety_text_moderation","hook_point":"output","endpoint":"{}","api_key":"k","stream_processing_mode":"buffer_full","max_buffer_bytes":4,"on_buffer_exceeded":"fail_open"}}"#, + acs.uri() + ); + let wire = run_textmod_stream(&cfg, &two_content_chunks("abcd", "efghij")).await; + assert!( + wire.contains("abcd"), + "fail-open cap must release held content; got:\n{wire}" + ); + assert!( + wire.contains("data: [DONE]"), + "fail-open released stream ends with [DONE]; got:\n{wire}" + ); + assert!( + !wire.contains("event: error"), + "fail-open release emits no error; got:\n{wire}" + ); + } + // ---- regression coverage for issue #107 ------------------------- // Pre-fix only /v1/chat/completions enforced rate-limit / budget; // every other LLM endpoint silently bypassed both. The test below diff --git a/schemas/resources/guardrail.schema.json b/schemas/resources/guardrail.schema.json index f4fba0b1..b83a6b73 100644 --- a/schemas/resources/guardrail.schema.json +++ b/schemas/resources/guardrail.schema.json @@ -106,6 +106,127 @@ "minimum": 0.0 } } + }, + { + "description": "Azure AI Content Safety Text Moderation. Category-severity + blocklist moderation via the `/contentsafety/text:analyze` API, on input and/or output (including streaming output). P2 (PRD-09c §6 P2, #379).", + "type": "object", + "required": [ + "api_key", + "endpoint", + "kind" + ], + "properties": { + "api_key": { + "description": "Subscription key (`Ocp-Apim-Subscription-Key`). Plaintext in memory only, never logged.", + "type": "string" + }, + "blocklist_names": { + "description": "Azure CS blocklist names to match against.", + "default": [], + "type": "array", + "items": { + "type": "string" + } + }, + "categories": { + "description": "Categories to analyze. Defaults to all four.", + "default": [ + "Hate", + "Sexual", + "SelfHarm", + "Violence" + ], + "type": "array", + "items": { + "type": "string" + } + }, + "endpoint": { + "description": "Azure Cognitive Services resource endpoint. The DP appends `/contentsafety/text:analyze?api-version=2024-09-01`.", + "type": "string" + }, + "halt_on_blocklist_hit": { + "description": "Forwarded to Azure's `haltOnBlocklistHit`.", + "default": false, + "type": "boolean" + }, + "kind": { + "type": "string", + "enum": [ + "azure_content_safety_text_moderation" + ] + }, + "max_buffer_bytes": { + "description": "Max bytes buffered in `buffer_full` mode before `on_buffer_exceeded` applies. Default 262 144.", + "default": 262144, + "type": "integer", + "format": "uint64", + "minimum": 0.0 + }, + "on_buffer_exceeded": { + "description": "`fail_closed` (default) or `fail_open` when the buffer cap is hit.", + "default": "fail_closed", + "type": "string" + }, + "output_fail_open": { + "description": "Fail-open policy for the OUTPUT hook. Defaults `false` (fail-closed) so an Azure outage can't release unscanned model output.", + "default": false, + "type": "boolean" + }, + "output_type": { + "description": "`FourSeverityLevels` (0,2,4,6; default) or `EightSeverityLevels` (0..7).", + "default": "FourSeverityLevels", + "type": "string" + }, + "severity_threshold": { + "description": "General severity threshold; a category at or above it blocks.", + "default": 2, + "type": "integer", + "format": "uint8", + "minimum": 0.0 + }, + "severity_threshold_by_category": { + "description": "Per-category threshold overrides (take precedence over the general one).", + "default": {}, + "type": "object", + "additionalProperties": { + "type": "integer", + "format": "uint8", + "minimum": 0.0 + } + }, + "stream_processing_mode": { + "description": "`window` (sliding-window incremental release; default) or `buffer_full` (whole-response hold-back).", + "default": "window", + "type": "string" + }, + "text_source": { + "description": "Input-hook text selection: `concatenate_user_content` (default) or `concatenate_all_content`. Ignored on the output hook.", + "default": "concatenate_user_content", + "type": "string" + }, + "timeout_ms": { + "description": "HTTP call timeout (ms); `fail_open` / `output_fail_open` govern the verdict when it elapses. See the `AzureContentSafetyConfig` note on why `0` means \"fire immediately\", not \"no timeout\".", + "default": 5000, + "type": "integer", + "format": "uint32", + "minimum": 0.0 + }, + "window_overlap_size": { + "description": "Chars carried between windows so a span split across a boundary is still caught. Default 256.", + "default": 256, + "type": "integer", + "format": "uint32", + "minimum": 0.0 + }, + "window_size": { + "description": "Sliding-window size in chars (window mode); cp-api caps it at the 10 000-char Azure limit. Default 10 000.", + "default": 10000, + "type": "integer", + "format": "uint32", + "minimum": 0.0 + } + } } ], "required": [