From 209346569993fe794001cb0d2be0ffbbd602d917 Mon Sep 17 00:00:00 2001 From: Jarvis Date: Thu, 2 Jul 2026 19:10:19 +0800 Subject: [PATCH 1/4] feat(bedrock): classify ApplyGuardrail response as allow/block/mask Bedrock reports action=GUARDRAIL_INTERVENED for both a hard block and a PII anonymization. Add classify_response() to tell them apart via the per-policy assessment actions (topic/content/word/contextual hits and PII/regex action=BLOCKED are hard blocks; PII action=ANONYMIZED with masked outputs is a mask), mirroring LiteLLM's block-iff-any-BLOCKED rule. Foundational only: apply() still blocks on a mask (fail-safe, behaviour- preserving) until the async rewrite channel that writes the masked outputs[].text back onto the request/response lands. Refs api7/AISIX-Cloud#932 --- crates/aisix-guardrails/src/bedrock.rs | 250 +++++++++++++++++++++++-- 1 file changed, 233 insertions(+), 17 deletions(-) diff --git a/crates/aisix-guardrails/src/bedrock.rs b/crates/aisix-guardrails/src/bedrock.rs index 8319f8e9..d1c264be 100644 --- a/crates/aisix-guardrails/src/bedrock.rs +++ b/crates/aisix-guardrails/src/bedrock.rs @@ -38,9 +38,10 @@ use async_trait::async_trait; use aws_credential_types::provider::SharedCredentialsProvider; use aws_credential_types::Credentials; use aws_sdk_bedrockruntime::config::{BehaviorVersion, Region}; -use aws_sdk_bedrockruntime::operation::apply_guardrail::ApplyGuardrailError; +use aws_sdk_bedrockruntime::operation::apply_guardrail::{ApplyGuardrailError, ApplyGuardrailOutput}; use aws_sdk_bedrockruntime::types::{ - GuardrailAction, GuardrailContentBlock, GuardrailContentSource, GuardrailTextBlock, + GuardrailAction, GuardrailAssessment, GuardrailContentBlock, GuardrailContentSource, + GuardrailSensitiveInformationPolicyAction, GuardrailTextBlock, }; use aws_sdk_bedrockruntime::Client; use aws_smithy_runtime_api::client::result::SdkError; @@ -205,24 +206,23 @@ impl BedrockGuardrail { }; match result { - Ok(resp) => match resp.action() { - GuardrailAction::GuardrailIntervened => GuardrailVerdict::block(format!( + Ok(resp) => match classify_response(&resp, &self.guardrail_id) { + BedrockOutcome::Allow => GuardrailVerdict::Allow, + BedrockOutcome::Block => GuardrailVerdict::block(format!( "bedrock guardrail {} intervened", self.guardrail_id )), - GuardrailAction::None => GuardrailVerdict::Allow, - other => { - // Forward-compat: an unknown enum variant from a - // future SDK upgrade. Treat as no-block (the - // safer interpretation since `intervened` is the - // active-block signal). - tracing::warn!( - guardrail_id = %self.guardrail_id, - action = ?other, - "unknown ApplyGuardrail action; treating as Allow", - ); - GuardrailVerdict::Allow - } + // Bedrock ANONYMIZED (masked) the content rather than + // blocking. The masked replacement text is in the payload. + // TODO(#932 bedrock follow-up): once the async rewrite + // channel lands, return a mask verdict and write the masked + // text back instead of blocking. Until then, block so the + // un-masked content is never released (fail-safe, and + // behaviour-preserving vs the pre-classify code). + BedrockOutcome::Mask(_) => GuardrailVerdict::block(format!( + "bedrock guardrail {} anonymized content", + self.guardrail_id + )), }, Err(failure) => self.handle_failure(failure, fail_open), } @@ -249,6 +249,97 @@ impl BedrockGuardrail { } } +/// The masking-aware interpretation of an `ApplyGuardrail` response. +/// +/// Bedrock reports `action = GUARDRAIL_INTERVENED` for BOTH a hard block +/// AND a PII anonymization (mask). The two are told apart by the +/// per-policy actions inside `assessments`: a topic/content/word/ +/// contextual-grounding policy hit, or a PII/regex entity with +/// `action = BLOCKED`, is a hard block; a PII/regex entity with +/// `action = ANONYMIZED` (and nothing blocking) is a mask, whose +/// replacement text Bedrock returns in `outputs[].text`. Mirrors +/// LiteLLM's `_should_raise_guardrail_blocked_exception` (raise iff any +/// assessment entry is BLOCKED; otherwise apply the masked output). +#[derive(Debug, PartialEq, Eq)] +enum BedrockOutcome { + /// `action = NONE` — nothing detected. + Allow, + /// A hard block: some policy blocked (topic/content/word/contextual) + /// or a PII/regex entity had `action = BLOCKED`. + Block, + /// Only anonymization occurred. Carries the masked replacement text + /// per `outputs[]` block (real Bedrock returns one). + Mask(Vec), +} + +/// Classify an `ApplyGuardrail` response into allow / block / mask. +/// Secure by default: an intervention that is neither a recognizable +/// block nor accompanied by masked output is treated as a block. +fn classify_response(resp: &ApplyGuardrailOutput, guardrail_id: &str) -> BedrockOutcome { + match resp.action() { + GuardrailAction::None => BedrockOutcome::Allow, + GuardrailAction::GuardrailIntervened => { + if resp.assessments().iter().any(assessment_has_hard_block) { + return BedrockOutcome::Block; + } + let masked: Vec = resp + .outputs() + .iter() + .filter_map(|o| o.text()) + .filter(|s| !s.is_empty()) + .map(str::to_owned) + .collect(); + if masked.is_empty() { + // Intervened, no recognizable hard block, no masked + // output — block rather than risk releasing content whose + // disposition we can't read. + BedrockOutcome::Block + } else { + BedrockOutcome::Mask(masked) + } + } + other => { + // Forward-compat: an unknown enum variant from a future SDK + // upgrade. `intervened` is the active signal, so an unknown + // action is treated as no-intervention (Allow). + tracing::warn!( + guardrail_id = %guardrail_id, + action = ?other, + "unknown ApplyGuardrail action; treating as Allow", + ); + BedrockOutcome::Allow + } + } +} + +/// True if the assessment carries any BLOCKING disposition — a +/// topic/content/word policy hit (these have no anonymize mode, so their +/// presence is a block), a contextual-grounding filter that BLOCKED, or a +/// PII/regex entity whose action is BLOCKED (as opposed to ANONYMIZED). +fn assessment_has_hard_block(a: &GuardrailAssessment) -> bool { + let topic_blocked = a.topic_policy().is_some_and(|p| !p.topics().is_empty()); + let content_blocked = a.content_policy().is_some_and(|p| !p.filters().is_empty()); + let word_blocked = a.word_policy().is_some_and(|p| { + !p.custom_words().is_empty() || !p.managed_word_lists().is_empty() + }); + let grounding_blocked = a.contextual_grounding_policy().is_some_and(|p| { + p.filters().iter().any(|f| { + *f.action() + == aws_sdk_bedrockruntime::types::GuardrailContextualGroundingPolicyAction::Blocked + }) + }); + let pii_blocked = a.sensitive_information_policy().is_some_and(|sip| { + sip.pii_entities() + .iter() + .any(|e| *e.action() == GuardrailSensitiveInformationPolicyAction::Blocked) + || sip + .regexes() + .iter() + .any(|r| *r.action() == GuardrailSensitiveInformationPolicyAction::Blocked) + }); + topic_blocked || content_blocked || word_blocked || grounding_blocked || pii_blocked +} + /// Failure cause buckets that map onto `guardrail_bypassed_reason` /// telemetry tags. `Other` collapses every long-tail SDK error onto /// `bedrock_5xx` so an unrecognised AWS error doesn't leak its @@ -380,6 +471,131 @@ mod tests { use super::*; use aisix_core::models::{BedrockAWSCredentials, BedrockConfig, BedrockLatencyMode}; + // --- classify_response: block vs mask vs allow (#932 bedrock) --------- + mod classify { + use super::super::{assessment_has_hard_block, classify_response, BedrockOutcome}; + use aws_sdk_bedrockruntime::operation::apply_guardrail::ApplyGuardrailOutput; + use aws_sdk_bedrockruntime::types::{ + GuardrailAction, GuardrailAssessment, GuardrailOutputContent, GuardrailPiiEntityFilter, + GuardrailPiiEntityType, GuardrailSensitiveInformationPolicyAction as PiiAction, + GuardrailSensitiveInformationPolicyAssessment, GuardrailTopic, GuardrailTopicPolicyAction, + GuardrailTopicPolicyAssessment, GuardrailTopicType, + }; + + fn resp( + action: GuardrailAction, + outputs: Vec<&str>, + assessments: Vec, + ) -> ApplyGuardrailOutput { + ApplyGuardrailOutput::builder() + .action(action) + .set_outputs(Some( + outputs + .into_iter() + .map(|t| GuardrailOutputContent::builder().text(t).build()) + .collect(), + )) + .set_assessments(Some(assessments)) + .build() + .expect("action/outputs/assessments all set") + } + + fn pii(action: PiiAction) -> GuardrailAssessment { + let entity = GuardrailPiiEntityFilter::builder() + .r#match("alice@example.com") + .r#type(GuardrailPiiEntityType::Email) + .action(action) + .build() + .unwrap(); + let sip = GuardrailSensitiveInformationPolicyAssessment::builder() + .pii_entities(entity) + .set_regexes(Some(vec![])) + .build() + .unwrap(); + GuardrailAssessment::builder() + .sensitive_information_policy(sip) + .build() + } + + fn topic() -> GuardrailAssessment { + let t = GuardrailTopic::builder() + .name("blocked-topic") + .r#type(GuardrailTopicType::Deny) + .action(GuardrailTopicPolicyAction::Blocked) + .build() + .unwrap(); + let tp = GuardrailTopicPolicyAssessment::builder() + .topics(t) + .build() + .unwrap(); + GuardrailAssessment::builder().topic_policy(tp).build() + } + + #[test] + fn action_none_is_allow() { + let r = resp(GuardrailAction::None, vec![], vec![]); + assert_eq!(classify_response(&r, "gid"), BedrockOutcome::Allow); + } + + #[test] + fn anonymized_pii_with_masked_output_is_mask() { + let r = resp( + GuardrailAction::GuardrailIntervened, + vec!["contact {EMAIL} about the order"], + vec![pii(PiiAction::Anonymized)], + ); + assert_eq!( + classify_response(&r, "gid"), + BedrockOutcome::Mask(vec!["contact {EMAIL} about the order".to_owned()]) + ); + } + + #[test] + fn blocked_pii_is_block() { + let r = resp( + GuardrailAction::GuardrailIntervened, + vec!["irrelevant"], + vec![pii(PiiAction::Blocked)], + ); + assert_eq!(classify_response(&r, "gid"), BedrockOutcome::Block); + } + + #[test] + fn topic_policy_hit_is_block_even_with_masked_output() { + // A hard block (topic) wins even if masked output is present. + let r = resp( + GuardrailAction::GuardrailIntervened, + vec!["masked text"], + vec![topic()], + ); + assert_eq!(classify_response(&r, "gid"), BedrockOutcome::Block); + } + + #[test] + fn mixed_anonymized_and_blocked_is_block() { + let r = resp( + GuardrailAction::GuardrailIntervened, + vec!["masked"], + vec![pii(PiiAction::Anonymized), pii(PiiAction::Blocked)], + ); + assert_eq!(classify_response(&r, "gid"), BedrockOutcome::Block); + } + + #[test] + fn intervened_without_hard_block_or_masked_output_is_block() { + // Secure default: an intervention we can't read as a mask blocks. + let r = resp(GuardrailAction::GuardrailIntervened, vec![], vec![]); + assert_eq!(classify_response(&r, "gid"), BedrockOutcome::Block); + } + + #[test] + fn hard_block_helper_only_true_for_blocking_dispositions() { + assert!(!assessment_has_hard_block(&pii(PiiAction::Anonymized))); + assert!(assessment_has_hard_block(&pii(PiiAction::Blocked))); + assert!(assessment_has_hard_block(&topic())); + } + } + fn cfg() -> BedrockConfig { BedrockConfig { guardrail_id: "abcdefgh1234".into(), From 9644c1e1c215372a99ff394d8b39eab874e11c72 Mon Sep 17 00:00:00 2001 From: Jarvis Date: Mon, 6 Jul 2026 12:27:52 +0800 Subject: [PATCH 2/4] =?UTF-8?q?feat(bedrock):=20honor=20ANONYMIZE=20?= =?UTF-8?q?=E2=80=94=20mask-and-continue=20via=20segment=20moderation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bedrock reports GUARDRAIL_INTERVENED for both block and mask; until now any intervention blocked. Now the four chat-shaped families (chat, messages, responses incl. bridge, completions) run a segment pass: the request/response text slots are sent as one content block each in a single ApplyGuardrail call (positional outputs[] contract), a hard-block disposition still blocks, and an ANONYMIZE disposition writes Bedrock's masked text back into the wire body and continues — LiteLLM's semantics, including its defensive fallback (misaligned outputs are never applied; originals stand). - Guardrail trait: moderates_segments + moderate_{input,output}_segments (SegmentsOutcome carries verdict + positional mask + entity counts) and check_{input,output}_non_segment so segment members are consulted exactly once per hook; GuardrailChain folds both passes with the usual attribution/bypass semantics and refuses drifted masks. - bedrock.rs: send() takes N content blocks; apply_segments() classifies block-vs-mask per assessments and aligns outputs[i]↔texts[i]; blob-mode check_* keeps mapping ANONYMIZE→Block for families with no write-back channel (embeddings/rerank/images/audio/passthrough/MCP unchanged). - proxy redact.rs: SegmentCollector/SegmentApplier reuse the #932 wire walkers for slot enumeration (collect→call→apply, same order by construction); moderate_body() folds the split verdicts; streaming hold-back paths rebuild the content-capture accumulator from the masked chunks/SSE so exporters never see pre-mask text (#947). - redacted_entity_counts gains Bedrock entity TYPES (EMAIL, …) — names only, matched values are never read (#153/#932). Part of the #932 Bedrock follow-up (AISIX-Cloud#932). --- crates/aisix-guardrails/src/bedrock.rs | 423 +++++++++++++++++++-- crates/aisix-guardrails/src/chain.rs | 274 ++++++++++++- crates/aisix-guardrails/src/lib.rs | 109 ++++++ crates/aisix-proxy/src/chat.rs | 86 ++++- crates/aisix-proxy/src/completions.rs | 32 +- crates/aisix-proxy/src/messages.rs | 139 ++++++- crates/aisix-proxy/src/redact.rs | 246 ++++++++++++ crates/aisix-proxy/src/responses.rs | 75 +++- crates/aisix-proxy/src/responses_bridge.rs | 52 ++- 9 files changed, 1373 insertions(+), 63 deletions(-) diff --git a/crates/aisix-guardrails/src/bedrock.rs b/crates/aisix-guardrails/src/bedrock.rs index d1c264be..ef558e54 100644 --- a/crates/aisix-guardrails/src/bedrock.rs +++ b/crates/aisix-guardrails/src/bedrock.rs @@ -16,7 +16,8 @@ //! | Bedrock response | `fail_open` | Verdict | //! |---------------------------------|-------------|--------------------------------| //! | `action=NONE` | n/a | Allow | -//! | `action=GUARDRAIL_INTERVENED` | n/a | Block { reason } | +//! | intervened, hard block | n/a | Block { reason } | +//! | intervened, ANONYMIZED only | n/a | masked write-back on the segment path (`moderate_*_segments`); Block on the blob path (`check_*`, no write-back channel) | //! | 5xx / IO error | true | Bypass { "bedrock_5xx" } | //! | 5xx / IO error | false | Block { "bedrock unavailable" } | //! | timeout (`latency_mode=timed`) | true | Bypass { "bedrock_timeout" } | @@ -38,7 +39,9 @@ use async_trait::async_trait; use aws_credential_types::provider::SharedCredentialsProvider; use aws_credential_types::Credentials; use aws_sdk_bedrockruntime::config::{BehaviorVersion, Region}; -use aws_sdk_bedrockruntime::operation::apply_guardrail::{ApplyGuardrailError, ApplyGuardrailOutput}; +use aws_sdk_bedrockruntime::operation::apply_guardrail::{ + ApplyGuardrailError, ApplyGuardrailOutput, +}; use aws_sdk_bedrockruntime::types::{ GuardrailAction, GuardrailAssessment, GuardrailContentBlock, GuardrailContentSource, GuardrailSensitiveInformationPolicyAction, GuardrailTextBlock, @@ -47,7 +50,7 @@ use aws_sdk_bedrockruntime::Client; use aws_smithy_runtime_api::client::result::SdkError; use aws_smithy_runtime_api::http::Response; -use crate::{Guardrail, GuardrailVerdict}; +use crate::{Guardrail, GuardrailVerdict, SegmentsOutcome}; /// One Bedrock guardrail row, materialised into a request-time /// dispatcher. Built once per snapshot from @@ -174,25 +177,30 @@ impl BedrockGuardrail { } } - /// Run `ApplyGuardrail` against a content block. Wraps the - /// SDK call with `latency_mode` enforcement and translates the - /// response/error into a `GuardrailVerdict` per §Behavior matrix. - async fn apply(&self, source: GuardrailContentSource, text: String) -> GuardrailVerdict { - let fail_open = self.fail_open_for(&source); - let req = self + /// One `ApplyGuardrail` call carrying `texts` as one content block + /// each (positional — `outputs[i]` aligns with `texts[i]` when + /// Bedrock anonymizes), wrapped with `latency_mode` enforcement. + async fn send( + &self, + source: GuardrailContentSource, + texts: &[String], + ) -> Result { + let mut req = self .client .apply_guardrail() .guardrail_identifier(&self.guardrail_id) .guardrail_version(&self.guardrail_version) - .source(source) - .content(GuardrailContentBlock::Text( + .source(source); + for text in texts { + req = req.content(GuardrailContentBlock::Text( GuardrailTextBlock::builder() .text(text) .build() .expect("GuardrailTextBlock requires text — set above"), )); + } - let result = match self.latency_mode { + match self.latency_mode { BedrockLatencyMode::Serial => req.send().await.map_err(BedrockFailure::from_sdk), BedrockLatencyMode::Timed { timeout_ms } => { match tokio::time::timeout(Duration::from_millis(timeout_ms as u64), req.send()) @@ -203,22 +211,23 @@ impl BedrockGuardrail { Err(_) => Err(BedrockFailure::Timeout), } } - }; + } + } - match result { + /// Blob-mode `ApplyGuardrail`: one joined content block, verdict only. + /// Serves `check_input`/`check_output` — the families with no mask + /// write-back channel — so an ANONYMIZE disposition maps to Block + /// there (releasing the un-masked content would defeat the operator's + /// policy; the segment path is where masking is honored). + async fn apply(&self, source: GuardrailContentSource, text: String) -> GuardrailVerdict { + let fail_open = self.fail_open_for(&source); + match self.send(source, std::slice::from_ref(&text)).await { Ok(resp) => match classify_response(&resp, &self.guardrail_id) { BedrockOutcome::Allow => GuardrailVerdict::Allow, BedrockOutcome::Block => GuardrailVerdict::block(format!( "bedrock guardrail {} intervened", self.guardrail_id )), - // Bedrock ANONYMIZED (masked) the content rather than - // blocking. The masked replacement text is in the payload. - // TODO(#932 bedrock follow-up): once the async rewrite - // channel lands, return a mask verdict and write the masked - // text back instead of blocking. Until then, block so the - // un-masked content is never released (fail-safe, and - // behaviour-preserving vs the pre-classify code). BedrockOutcome::Mask(_) => GuardrailVerdict::block(format!( "bedrock guardrail {} anonymized content", self.guardrail_id @@ -228,6 +237,49 @@ impl BedrockGuardrail { } } + /// Segment-mode `ApplyGuardrail`: one content block per text slot, + /// verdict + positional mask write-back. On an ANONYMIZE disposition + /// Bedrock returns one `outputs[]` entry per input block; when that + /// alignment holds the masked texts are returned for write-back. + /// When it doesn't (a provider quirk we can't attribute to slots), + /// keep the originals and continue — LiteLLM's `_merge_masked_texts` + /// fallback: never misapply masked content to the wrong slot. + async fn apply_segments( + &self, + source: GuardrailContentSource, + texts: &[String], + ) -> SegmentsOutcome { + let fail_open = self.fail_open_for(&source); + match self.send(source, texts).await { + Ok(resp) => match classify_response(&resp, &self.guardrail_id) { + BedrockOutcome::Allow => SegmentsOutcome::allow(), + BedrockOutcome::Block => SegmentsOutcome::from_verdict(GuardrailVerdict::block( + format!("bedrock guardrail {} intervened", self.guardrail_id), + )), + BedrockOutcome::Mask(outputs) => { + if outputs.len() == texts.len() { + SegmentsOutcome { + verdict: GuardrailVerdict::Allow, + masked: Some(outputs), + counts: anonymized_counts(&resp), + } + } else { + tracing::warn!( + row = %self.row_name, + guardrail_id = %self.guardrail_id, + expected = texts.len(), + got = outputs.len(), + "bedrock masked outputs don't align with input \ + blocks; skipping mask write-back", + ); + SegmentsOutcome::allow() + } + } + }, + Err(failure) => SegmentsOutcome::from_verdict(self.handle_failure(failure, fail_open)), + } + } + fn handle_failure(&self, failure: BedrockFailure, fail_open: bool) -> GuardrailVerdict { let (reason, error_detail, error_source) = failure.log_fields(); tracing::warn!( @@ -268,7 +320,9 @@ enum BedrockOutcome { /// or a PII/regex entity had `action = BLOCKED`. Block, /// Only anonymization occurred. Carries the masked replacement text - /// per `outputs[]` block (real Bedrock returns one). + /// per `outputs[]` block, in order and WITHOUT dropping empty entries + /// — `outputs[i]` must keep aligning with the i-th input content + /// block for the segment write-back. Mask(Vec), } @@ -285,11 +339,9 @@ fn classify_response(resp: &ApplyGuardrailOutput, guardrail_id: &str) -> Bedrock let masked: Vec = resp .outputs() .iter() - .filter_map(|o| o.text()) - .filter(|s| !s.is_empty()) - .map(str::to_owned) + .map(|o| o.text().unwrap_or_default().to_owned()) .collect(); - if masked.is_empty() { + if masked.iter().all(String::is_empty) { // Intervened, no recognizable hard block, no masked // output — block rather than risk releasing content whose // disposition we can't read. @@ -312,6 +364,33 @@ fn classify_response(resp: &ApplyGuardrailOutput, guardrail_id: &str) -> Bedrock } } +/// Per-entity counts of what Bedrock ANONYMIZED, for +/// `redacted_entity_counts` telemetry. Keys are the PII entity TYPE +/// (`EMAIL`, `PHONE`, …) or the operator's configured regex name — +/// config-level metadata, never matched values, so the map is safe to +/// log and attach to telemetry (#153 / #932 no-leak criterion). The +/// assessment's `match` fields are deliberately never read. +fn anonymized_counts(resp: &ApplyGuardrailOutput) -> std::collections::BTreeMap { + let mut counts = std::collections::BTreeMap::new(); + for a in resp.assessments() { + let Some(sip) = a.sensitive_information_policy() else { + continue; + }; + for e in sip.pii_entities() { + if *e.action() == GuardrailSensitiveInformationPolicyAction::Anonymized { + *counts.entry(e.r#type().as_str().to_owned()).or_insert(0) += 1; + } + } + for r in sip.regexes() { + if *r.action() == GuardrailSensitiveInformationPolicyAction::Anonymized { + let name = r.name().unwrap_or("regex").to_owned(); + *counts.entry(name).or_insert(0) += 1; + } + } + } + counts +} + /// True if the assessment carries any BLOCKING disposition — a /// topic/content/word policy hit (these have no anonymize mode, so their /// presence is a block), a contextual-grounding filter that BLOCKED, or a @@ -319,9 +398,9 @@ fn classify_response(resp: &ApplyGuardrailOutput, guardrail_id: &str) -> Bedrock fn assessment_has_hard_block(a: &GuardrailAssessment) -> bool { let topic_blocked = a.topic_policy().is_some_and(|p| !p.topics().is_empty()); let content_blocked = a.content_policy().is_some_and(|p| !p.filters().is_empty()); - let word_blocked = a.word_policy().is_some_and(|p| { - !p.custom_words().is_empty() || !p.managed_word_lists().is_empty() - }); + let word_blocked = a + .word_policy() + .is_some_and(|p| !p.custom_words().is_empty() || !p.managed_word_lists().is_empty()); let grounding_blocked = a.contextual_grounding_policy().is_some_and(|p| { p.filters().iter().any(|f| { *f.action() @@ -451,6 +530,42 @@ impl Guardrail for BedrockGuardrail { } self.apply(GuardrailContentSource::Output, text).await } + + /// Bedrock moderates via the segment pass on call sites that support + /// mask write-back; those sites pair `moderate_*_segments` with + /// `check_*_non_segment`, so the guardrail is called exactly once. + fn moderates_segments(&self) -> bool { + true + } + + async fn moderate_input_segments(&self, texts: &[String]) -> SegmentsOutcome { + if !matches!( + self.hook_point, + GuardrailHookPoint::Input | GuardrailHookPoint::Both + ) { + return SegmentsOutcome::allow(); + } + if texts.iter().all(|t| t.is_empty()) { + // Nothing to scan — Bedrock would 400 on empty content. + return SegmentsOutcome::allow(); + } + self.apply_segments(GuardrailContentSource::Input, texts) + .await + } + + async fn moderate_output_segments(&self, texts: &[String]) -> SegmentsOutcome { + if !matches!( + self.hook_point, + GuardrailHookPoint::Output | GuardrailHookPoint::Both + ) { + return SegmentsOutcome::allow(); + } + if texts.iter().all(|t| t.is_empty()) { + return SegmentsOutcome::allow(); + } + self.apply_segments(GuardrailContentSource::Output, texts) + .await + } } /// Concatenate the request's user-visible message contents into one @@ -478,8 +593,8 @@ mod tests { use aws_sdk_bedrockruntime::types::{ GuardrailAction, GuardrailAssessment, GuardrailOutputContent, GuardrailPiiEntityFilter, GuardrailPiiEntityType, GuardrailSensitiveInformationPolicyAction as PiiAction, - GuardrailSensitiveInformationPolicyAssessment, GuardrailTopic, GuardrailTopicPolicyAction, - GuardrailTopicPolicyAssessment, GuardrailTopicType, + GuardrailSensitiveInformationPolicyAssessment, GuardrailTopic, + GuardrailTopicPolicyAction, GuardrailTopicPolicyAssessment, GuardrailTopicType, }; fn resp( @@ -594,6 +709,52 @@ mod tests { assert!(assessment_has_hard_block(&pii(PiiAction::Blocked))); assert!(assessment_has_hard_block(&topic())); } + + /// Positional integrity: an empty `outputs[]` entry is preserved, + /// not dropped — `outputs[i]` must keep aligning with the i-th + /// input content block for the segment write-back. + #[test] + fn mask_preserves_empty_output_positions() { + let r = resp( + GuardrailAction::GuardrailIntervened, + vec!["", "masked tail"], + vec![pii(PiiAction::Anonymized)], + ); + assert_eq!( + classify_response(&r, "gid"), + BedrockOutcome::Mask(vec![String::new(), "masked tail".to_owned()]) + ); + } + + /// The anonymized-entity counts read TYPE/name metadata only — + /// the matched value ("alice@example.com" in the fixture) never + /// appears in a key. + #[test] + fn anonymized_counts_carry_types_not_values() { + let r = resp( + GuardrailAction::GuardrailIntervened, + vec!["{EMAIL}"], + vec![pii(PiiAction::Anonymized), pii(PiiAction::Anonymized)], + ); + let counts = super::super::anonymized_counts(&r); + assert_eq!(counts.get("EMAIL"), Some(&2)); + assert_eq!(counts.len(), 1); + assert!( + !counts.keys().any(|k| k.contains("alice")), + "matched values must never leak into count keys", + ); + } + + /// BLOCKED entities don't show up in the anonymize counts. + #[test] + fn anonymized_counts_skip_blocked_entities() { + let r = resp( + GuardrailAction::GuardrailIntervened, + vec![], + vec![pii(PiiAction::Blocked)], + ); + assert!(super::super::anonymized_counts(&r).is_empty()); + } } fn cfg() -> BedrockConfig { @@ -1009,4 +1170,204 @@ mod tests { other => panic!("expected Bypass(bedrock_timeout), got {other:?}"), } } + + // --- segment mode (mask write-back, #932 bedrock follow-up) ---------- + + fn anonymized_body(outputs: Vec<&str>) -> serde_json::Value { + json!({ + "action": "GUARDRAIL_INTERVENED", + "outputs": outputs.into_iter().map(|t| json!({"text": t})).collect::>(), + "assessments": [{ + "sensitiveInformationPolicy": { + "piiEntities": [ + {"match": "alice@example.com", "type": "EMAIL", "action": "ANONYMIZED"} + ], + "regexes": [] + } + }], + "usage": { + "topicPolicyUnits": 0, + "contentPolicyUnits": 0, + "wordPolicyUnits": 0, + "sensitiveInformationPolicyUnits": 1, + "sensitiveInformationPolicyFreeUnits": 0, + "contextualGroundingPolicyUnits": 0 + } + }) + } + + /// Segment mode sends ONE call with one content block per text slot, + /// and an aligned ANONYMIZED response comes back as positional masked + /// texts + entity-type counts. + #[tokio::test] + async fn apply_segments_sends_one_block_per_text_and_masks_positionally() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path_regex(r"^/guardrail/.+/apply$")) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(anonymized_body(vec!["seg one", "mail {EMAIL} now"])), + ) + .expect(1) + .mount(&server) + .await; + + let g = build_with_endpoint(server.uri(), true); + let texts = vec![ + "seg one".to_owned(), + "mail alice@example.com now".to_owned(), + ]; + let outcome = g + .apply_segments(GuardrailContentSource::Input, &texts) + .await; + + assert_eq!(outcome.verdict, GuardrailVerdict::Allow); + assert_eq!( + outcome.masked, + Some(vec!["seg one".to_owned(), "mail {EMAIL} now".to_owned()]), + ); + assert_eq!(outcome.counts.get("EMAIL"), Some(&1)); + + // The single request must carry the two slots as two content + // blocks (positional contract with `outputs[]`). + let reqs = server.received_requests().await.unwrap(); + assert_eq!(reqs.len(), 1); + let body: serde_json::Value = serde_json::from_slice(&reqs[0].body).unwrap(); + let content = body.get("content").and_then(|c| c.as_array()).unwrap(); + assert_eq!(content.len(), 2, "one content block per text slot"); + assert_eq!( + content[0].pointer("/text/text").and_then(|v| v.as_str()), + Some("seg one"), + ); + assert_eq!( + content[1].pointer("/text/text").and_then(|v| v.as_str()), + Some("mail alice@example.com now"), + ); + } + + /// The defensive fallback (LiteLLM `_merge_masked_texts` semantics): + /// masked outputs that can't be aligned positionally are NOT applied + /// — originals stand, the request continues. + #[tokio::test] + async fn apply_segments_misaligned_outputs_keep_originals_and_allow() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path_regex(r"^/guardrail/.+/apply$")) + .respond_with( + // 2 input slots, only 1 output — cannot attribute. + ResponseTemplate::new(200).set_body_json(anonymized_body(vec!["one blob"])), + ) + .mount(&server) + .await; + + let g = build_with_endpoint(server.uri(), true); + let texts = vec!["a".to_owned(), "b".to_owned()]; + let outcome = g + .apply_segments(GuardrailContentSource::Input, &texts) + .await; + assert_eq!(outcome.verdict, GuardrailVerdict::Allow); + assert_eq!(outcome.masked, None, "misaligned mask must not be applied"); + } + + /// A hard block (BLOCKED PII entity) on the segment path still + /// blocks — masking never bypasses a blocking disposition. + #[tokio::test] + async fn apply_segments_hard_block_still_blocks() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path_regex(r"^/guardrail/.+/apply$")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "blocked message"}], + "assessments": [{ + "sensitiveInformationPolicy": { + "piiEntities": [ + {"match": "x", "type": "EMAIL", "action": "BLOCKED"} + ], + "regexes": [] + } + }], + "usage": { + "topicPolicyUnits": 0, + "contentPolicyUnits": 0, + "wordPolicyUnits": 0, + "sensitiveInformationPolicyUnits": 1, + "sensitiveInformationPolicyFreeUnits": 0, + "contextualGroundingPolicyUnits": 0 + } + }))) + .mount(&server) + .await; + + let g = build_with_endpoint(server.uri(), true); + let outcome = g + .apply_segments(GuardrailContentSource::Input, &["x".to_owned()]) + .await; + assert!(outcome.verdict.is_block()); + assert_eq!(outcome.masked, None); + } + + /// Blob mode (`check_*`, the families with no write-back channel) + /// keeps mapping an ANONYMIZE disposition to Block — releasing the + /// un-masked content there would defeat the operator's policy. + #[tokio::test] + async fn blob_apply_still_blocks_on_anonymize() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path_regex(r"^/guardrail/.+/apply$")) + .respond_with( + ResponseTemplate::new(200).set_body_json(anonymized_body(vec!["{EMAIL}"])), + ) + .mount(&server) + .await; + + let g = build_with_endpoint(server.uri(), true); + let v = g + .apply(GuardrailContentSource::Input, "alice@example.com".into()) + .await; + match v { + GuardrailVerdict::Block { reason, .. } => { + assert!(reason.contains("anonymized content"), "got {reason}"); + } + other => panic!("expected Block, got {other:?}"), + } + } + + /// Segment-path failures keep the fail-open/closed contract: a 5xx + /// with `fail_open=false` blocks, with `fail_open=true` bypasses. + #[tokio::test] + async fn apply_segments_5xx_honors_fail_open() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path_regex(r"^/guardrail/.+/apply$")) + .respond_with(ResponseTemplate::new(500).set_body_json(json!({ + "__type": "InternalServerException", + "message": "boom" + }))) + .mount(&server) + .await; + + let open = build_with_endpoint(server.uri(), true); + let outcome = open + .apply_segments(GuardrailContentSource::Input, &["x".to_owned()]) + .await; + assert!(outcome.verdict.is_bypass()); + + let closed = build_with_endpoint(server.uri(), false); + let outcome = closed + .apply_segments(GuardrailContentSource::Input, &["x".to_owned()]) + .await; + assert!(outcome.verdict.is_block()); + } + + /// Hook-point gating carries over to the segment hooks: an + /// Output-only row must not scan input segments (Allow without ever + /// hitting AWS), mirroring `output_only_row_skips_input_check`. + #[tokio::test] + async fn output_only_row_skips_input_segments() { + let mut g = build_test(true); + g.hook_point = GuardrailHookPoint::Output; + let outcome = g.moderate_input_segments(&["hello".to_owned()]).await; + assert_eq!(outcome, SegmentsOutcome::allow()); + } } diff --git a/crates/aisix-guardrails/src/chain.rs b/crates/aisix-guardrails/src/chain.rs index f7414171..cf46bb0b 100644 --- a/crates/aisix-guardrails/src/chain.rs +++ b/crates/aisix-guardrails/src/chain.rs @@ -13,7 +13,7 @@ use aisix_core::AppliedGuardrail; use aisix_gateway::{ChatFormat, ChatResponse}; use async_trait::async_trait; -use crate::{Guardrail, GuardrailVerdict, Redaction, StreamOutputPolicy}; +use crate::{Guardrail, GuardrailVerdict, Redaction, SegmentsOutcome, StreamOutputPolicy}; /// One chain member: the runtime guardrail plus the operator-facing name /// of the row it was built from. The name is what `Block` verdicts are @@ -205,6 +205,73 @@ impl Guardrail for GuardrailChain { } } + fn moderates_segments(&self) -> bool { + self.members + .iter() + .any(|m| m.guardrail.moderates_segments()) + } + + /// Fold over segment-moderating members only. A Block short-circuits + /// (attributed like the check folds); masked texts COMPOSE — each + /// member moderates the previous member's masked output, mirroring + /// `fold_redactions`; the first Bypass reason sticks. Counts merge. + async fn moderate_input_segments(&self, texts: &[String]) -> SegmentsOutcome { + fold_segments(&self.members, texts, true).await + } + + async fn moderate_output_segments(&self, texts: &[String]) -> SegmentsOutcome { + fold_segments(&self.members, texts, false).await + } + + /// The check fold minus segment-moderating members — the pass those + /// members are consulted through is `moderate_*_segments`, run by the + /// same call sites. Recurses via the member's own + /// `check_input_non_segment` so a nested chain filters its own members + /// rather than being skipped wholesale. + async fn check_input_non_segment(&self, req: &ChatFormat) -> GuardrailVerdict { + let mut bypass: Option = None; + for m in &self.members { + match m.guardrail.check_input_non_segment(req).await { + GuardrailVerdict::Allow => continue, + GuardrailVerdict::Block { + reason, + guardrail_name, + } => return attribute_block(&m.name, reason, guardrail_name), + GuardrailVerdict::Bypass { reason } => { + if bypass.is_none() { + bypass = Some(reason); + } + } + } + } + match bypass { + Some(reason) => GuardrailVerdict::Bypass { reason }, + None => GuardrailVerdict::Allow, + } + } + + async fn check_output_non_segment(&self, resp: &ChatResponse) -> GuardrailVerdict { + let mut bypass: Option = None; + for m in &self.members { + match m.guardrail.check_output_non_segment(resp).await { + GuardrailVerdict::Allow => continue, + GuardrailVerdict::Block { + reason, + guardrail_name, + } => return attribute_block(&m.name, reason, guardrail_name), + GuardrailVerdict::Bypass { reason } => { + if bypass.is_none() { + bypass = Some(reason); + } + } + } + } + match bypass { + Some(reason) => GuardrailVerdict::Bypass { reason }, + None => GuardrailVerdict::Allow, + } + } + fn redacts_input(&self) -> bool { self.members.iter().any(|m| m.guardrail.redacts_input()) } @@ -237,6 +304,68 @@ impl Guardrail for GuardrailChain { } } +/// Fold the texts through each segment-moderating member. Mirrors the +/// check folds (first Block short-circuits with attribution, first Bypass +/// reason sticks) plus mask composition: each member moderates the +/// previous member's masked output. Counts merge across members. +async fn fold_segments(members: &[ChainMember], texts: &[String], input: bool) -> SegmentsOutcome { + let mut masked: Option> = None; + let mut counts = std::collections::BTreeMap::new(); + let mut bypass: Option = None; + for m in members { + if !m.guardrail.moderates_segments() { + continue; + } + let src: &[String] = masked.as_deref().unwrap_or(texts); + let outcome = if input { + m.guardrail.moderate_input_segments(src).await + } else { + m.guardrail.moderate_output_segments(src).await + }; + match outcome.verdict { + GuardrailVerdict::Allow => {} + GuardrailVerdict::Block { + reason, + guardrail_name, + } => { + return SegmentsOutcome::from_verdict(attribute_block( + &m.name, + reason, + guardrail_name, + )) + } + GuardrailVerdict::Bypass { reason } => { + if bypass.is_none() { + bypass = Some(reason); + } + } + } + if let Some(new_masked) = outcome.masked { + // Implementations uphold alignment with THEIR input; refuse a + // drifted length here so a broken member can't desync slots. + if new_masked.len() == src.len() { + masked = Some(new_masked); + } else { + tracing::warn!( + member = %m.name, + expected = src.len(), + got = new_masked.len(), + "segment moderation returned misaligned mask; keeping originals", + ); + } + } + Redaction::merge_counts(&mut counts, &outcome.counts); + } + SegmentsOutcome { + verdict: match bypass { + Some(reason) => GuardrailVerdict::Bypass { reason }, + None => GuardrailVerdict::Allow, + }, + masked, + counts, + } +} + /// Fold `text` through each member's redactor, merging counts. `None` /// when no member changed anything. fn fold_redactions<'a>( @@ -522,6 +651,149 @@ mod tests { assert!(!empty.stream_output_policy().holds_back()); } + // --- segment moderation folds (#932 bedrock follow-up) --------------- + + /// A stub segment moderator: uppercases every slot and reports a + /// fixed count key, or blocks/bypasses on demand. + struct StubSegments { + verdict: GuardrailVerdict, + mask: bool, + } + #[async_trait] + impl Guardrail for StubSegments { + fn name(&self) -> &'static str { + "stub-segments" + } + fn moderates_segments(&self) -> bool { + true + } + async fn check_input(&self, _req: &ChatFormat) -> GuardrailVerdict { + panic!("segment member must not be consulted via check_input_non_segment"); + } + async fn moderate_input_segments(&self, texts: &[String]) -> crate::SegmentsOutcome { + let mut counts = std::collections::BTreeMap::new(); + counts.insert("STUB".to_owned(), texts.len() as u32); + crate::SegmentsOutcome { + verdict: self.verdict.clone(), + masked: self + .mask + .then(|| texts.iter().map(|t| t.to_uppercase()).collect()), + counts, + } + } + } + + /// The non-segment check fold skips segment members (they're consulted + /// via the segment pass) while normal members still run — the panic in + /// the stub's `check_input` proves the skip. + #[tokio::test] + async fn check_input_non_segment_skips_segment_members_but_not_others() { + let chain = GuardrailChain::new(vec![ + Arc::new(StubSegments { + verdict: GuardrailVerdict::Allow, + mask: false, + }), + Arc::new(KeywordBlocklist::new(vec![KeywordRule::literal("AKIA")])), + ]); + // Keyword member still blocks... + assert!(chain + .check_input_non_segment(&req("here is AKIAEXAMPLE")) + .await + .is_block()); + // ...and a clean request is Allow (the stub's check_input would + // have panicked if consulted). + assert_eq!( + chain.check_input_non_segment(&req("clean")).await, + GuardrailVerdict::Allow, + ); + // The FULL fold still consults every member (unconverted call + // sites keep blob-mode coverage) — the stub panics to prove it + // WOULD be consulted there; assert via catch_unwind-free route: + // moderates_segments visibility. + assert!(chain.moderates_segments()); + } + + /// Segment masks compose across members in chain order, counts merge, + /// and a Block short-circuits with attribution. + #[tokio::test] + async fn segment_fold_composes_masks_and_attributes_blocks() { + // Two maskers: uppercase then uppercase again (idempotent — the + // composition is observable via counts merging to 2 members). + let chain = GuardrailChain::new_with_applied( + vec![ + ( + "mask-a".to_owned(), + Arc::new(StubSegments { + verdict: GuardrailVerdict::Allow, + mask: true, + }) as Arc, + ), + ( + "mask-b".to_owned(), + Arc::new(StubSegments { + verdict: GuardrailVerdict::Allow, + mask: true, + }), + ), + ], + Vec::new(), + ); + let texts = vec!["hello".to_owned(), "world".to_owned()]; + let out = chain.moderate_input_segments(&texts).await; + assert_eq!(out.verdict, GuardrailVerdict::Allow); + assert_eq!( + out.masked, + Some(vec!["HELLO".to_owned(), "WORLD".to_owned()]), + ); + assert_eq!(out.counts.get("STUB"), Some(&4), "2 members × 2 slots"); + + // Block short-circuits and is attributed to the firing member. + let blocking = GuardrailChain::new_with_applied( + vec![( + "seg-blocker".to_owned(), + Arc::new(StubSegments { + verdict: GuardrailVerdict::block("pii blocked"), + mask: false, + }) as Arc, + )], + Vec::new(), + ); + match blocking.moderate_input_segments(&texts).await.verdict { + GuardrailVerdict::Block { guardrail_name, .. } => { + assert_eq!(guardrail_name.as_deref(), Some("seg-blocker")) + } + other => panic!("expected Block, got {other:?}"), + } + } + + /// A member returning a mask whose length drifted from ITS input is + /// refused (originals kept) — the chain-level alignment guard. + #[tokio::test] + async fn segment_fold_refuses_misaligned_member_mask() { + struct Drifting; + #[async_trait] + impl Guardrail for Drifting { + fn name(&self) -> &'static str { + "drifting" + } + fn moderates_segments(&self) -> bool { + true + } + async fn moderate_input_segments(&self, _texts: &[String]) -> crate::SegmentsOutcome { + crate::SegmentsOutcome { + verdict: GuardrailVerdict::Allow, + masked: Some(vec!["only-one".to_owned()]), + counts: std::collections::BTreeMap::new(), + } + } + } + let chain = GuardrailChain::new(vec![Arc::new(Drifting)]); + let texts = vec!["a".to_owned(), "b".to_owned()]; + let out = chain.moderate_input_segments(&texts).await; + assert_eq!(out.masked, None, "drifted mask must be refused"); + assert_eq!(out.verdict, GuardrailVerdict::Allow); + } + #[test] fn new_has_empty_applied_and_new_with_applied_reports_it() { // `new` (the in-memory/test constructor) carries no applied metadata; diff --git a/crates/aisix-guardrails/src/lib.rs b/crates/aisix-guardrails/src/lib.rs index 44ed363c..a302f2ce 100644 --- a/crates/aisix-guardrails/src/lib.rs +++ b/crates/aisix-guardrails/src/lib.rs @@ -159,6 +159,20 @@ impl GuardrailVerdict { _ => None, } } + + /// Fold the verdicts of two split moderation passes over the same + /// content (the non-segment check + the segment pass) into one: + /// Block wins (`self` first), then Bypass (`self`'s reason first), + /// else Allow. + pub fn merged_with(self, other: GuardrailVerdict) -> GuardrailVerdict { + match (self, other) { + (b @ GuardrailVerdict::Block { .. }, _) => b, + (_, b @ GuardrailVerdict::Block { .. }) => b, + (by @ GuardrailVerdict::Bypass { .. }, _) => by, + (_, by @ GuardrailVerdict::Bypass { .. }) => by, + _ => GuardrailVerdict::Allow, + } + } } /// How a guardrail wants STREAMED output moderated. The proxy's SSE @@ -279,6 +293,46 @@ impl Redaction { } } +/// Outcome of [`Guardrail::moderate_input_segments`] / +/// [`Guardrail::moderate_output_segments`] — remote moderation of a +/// request's text segments in ONE provider call (kind=bedrock). +/// +/// `masked`, when present, is positionally aligned with the input +/// `texts` slice: `masked[i]` replaces `texts[i]`. Implementations MUST +/// uphold that alignment or return `masked: None` (the caller then keeps +/// the originals — the LiteLLM `_merge_masked_texts` defensive fallback: +/// never misapply masked content to the wrong slot). +/// +/// `counts` mirrors [`Redaction::counts`]: entity NAMES only (e.g. a +/// Bedrock PII entity type like `EMAIL`), never matched values, so it is +/// safe for logs and telemetry (#153 / #932 no-leak criterion). +#[derive(Debug, Clone, PartialEq)] +pub struct SegmentsOutcome { + pub verdict: GuardrailVerdict, + pub masked: Option>, + pub counts: std::collections::BTreeMap, +} + +impl SegmentsOutcome { + /// Plain Allow: nothing detected, nothing rewritten. + pub fn allow() -> Self { + Self { + verdict: GuardrailVerdict::Allow, + masked: None, + counts: std::collections::BTreeMap::new(), + } + } + + /// Wrap a bare verdict (Block/Bypass paths carry no mask or counts). + pub fn from_verdict(verdict: GuardrailVerdict) -> Self { + Self { + verdict, + masked: None, + counts: std::collections::BTreeMap::new(), + } + } +} + /// Pluggable content-policy hook. Production wires `Arc` /// in `ProxyState`; tests construct in-memory chains directly. #[async_trait] @@ -360,6 +414,61 @@ pub trait Guardrail: Send + Sync + 'static { fn redact_output_text(&self, _text: &str) -> Option { None } + + // --- remote segment moderation (#932 bedrock follow-up) --------------- + // + // A remote-API guardrail that can MASK (Bedrock PII anonymize) can't + // implement the sync per-field redact contract above — the mask comes + // back from the provider call itself. Instead the proxy hands such a + // guardrail ALL of a request's text segments at once (in wire-walker + // order), gets verdict + positionally-aligned masked replacements from + // ONE provider call, and writes them back per wire shape. Call sites + // that run this pass pair it with `check_*_non_segment` so the + // guardrail is consulted exactly once per hook. + + /// `true` when this guardrail moderates via the segment hooks below. + /// Such a member is skipped by `check_input_non_segment` / + /// `check_output_non_segment` (the segment pass covers it). + fn moderates_segments(&self) -> bool { + false + } + + /// Moderate the request's text segments in one remote call. Only + /// meaningful when [`Self::moderates_segments`] is `true`; the default + /// allows so a caller that runs the pass unconditionally is safe. + async fn moderate_input_segments(&self, _texts: &[String]) -> SegmentsOutcome { + SegmentsOutcome::allow() + } + + /// Moderate the response's text segments in one remote call. + async fn moderate_output_segments(&self, _texts: &[String]) -> SegmentsOutcome { + SegmentsOutcome::allow() + } + + /// `check_input` minus segment-moderating members — used by call + /// sites that ALSO run [`Self::moderate_input_segments`], so a + /// segment member isn't consulted twice (and billed twice). For a + /// leaf guardrail this is all-or-nothing: a segment moderator + /// answers via the segment pass (Allow here), anything else answers + /// via its normal check. [`GuardrailChain`] overrides with a + /// member-filtered fold. + async fn check_input_non_segment(&self, req: &ChatFormat) -> GuardrailVerdict { + if self.moderates_segments() { + GuardrailVerdict::Allow + } else { + self.check_input(req).await + } + } + + /// `check_output` minus segment-moderating members (see + /// [`Self::check_input_non_segment`]). + async fn check_output_non_segment(&self, resp: &ChatResponse) -> GuardrailVerdict { + if self.moderates_segments() { + GuardrailVerdict::Allow + } else { + self.check_output(resp).await + } + } } #[cfg(test)] diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index 63b939bd..f14bf7f2 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -863,7 +863,19 @@ async fn dispatch( // + fail_open=true) doesn't short-circuit; the reason is stashed // and attached to the telemetry event when the request finishes. let mut bypass_reason: Option = None; - match resolved_chain.check_input(req).await { + // Split moderation: local/blob members check first (on the original + // text), then the segment pass runs the Bedrock call and writes + // ANONYMIZE masks back into `req` (#932 bedrock follow-up). + let input_verdict = resolved_chain.check_input_non_segment(req).await; + let input_verdict = crate::redact::moderate_body( + resolved_chain.as_ref(), + crate::redact::Direction::Input, + input_verdict, + redactions_out, + |g| crate::redact::redact_chat_format(g, req), + ) + .await; + match input_verdict { GuardrailVerdict::Allow => {} GuardrailVerdict::Block { reason, @@ -1568,7 +1580,16 @@ async fn dispatch( // #448: a cache hit is client-visible output just like a // fresh upstream response, so it must run output guardrails // before being returned — not bypass them. - match resolved_chain.check_output(&cached).await { + let cached_verdict = resolved_chain.check_output_non_segment(&cached).await; + let cached_verdict = crate::redact::moderate_body( + resolved_chain.as_ref(), + crate::redact::Direction::Output, + cached_verdict, + redactions_out, + |g| crate::redact::redact_chat_response(g, &mut cached), + ) + .await; + match cached_verdict { GuardrailVerdict::Block { reason, guardrail_name, @@ -1913,7 +1934,16 @@ async fn dispatch( // ingesting telemetry; the DP just records 0.0 on the wire. let cost_usd = 0.0; - match resolved_chain.check_output(&upstream).await { + let output_verdict = resolved_chain.check_output_non_segment(&upstream).await; + let output_verdict = crate::redact::moderate_body( + resolved_chain.as_ref(), + crate::redact::Direction::Output, + output_verdict, + redactions_out, + |g| crate::redact::redact_chat_response(g, &mut upstream), + ) + .await; + match output_verdict { GuardrailVerdict::Allow => {} GuardrailVerdict::Block { reason, @@ -2702,7 +2732,19 @@ async fn dispatch_ensemble( // bills the full panel + judge): on a block we therefore pass // `charge: None` rather than a judge-only `UpstreamCharge`, which // would double-count the judge AND still miss the panel. - match resolved_chain.check_output(&outcome.response).await { + let mut ensemble_redactions = input_redactions.clone(); + let ensemble_verdict = resolved_chain + .check_output_non_segment(&outcome.response) + .await; + let ensemble_verdict = crate::redact::moderate_body( + resolved_chain.as_ref(), + crate::redact::Direction::Output, + ensemble_verdict, + &mut ensemble_redactions, + |g| crate::redact::redact_chat_response(g, &mut outcome.response), + ) + .await; + match ensemble_verdict { GuardrailVerdict::Allow => {} GuardrailVerdict::Block { reason, @@ -2742,7 +2784,6 @@ async fn dispatch_ensemble( // the check above ran on the original text, the client gets the masked // one — then emit the (non-blocked) sub-call events with the final // bypass value (which an output bypass above may have set). - let mut ensemble_redactions = input_redactions.clone(); crate::redact::merge_counts( &mut ensemble_redactions, crate::redact::redact_chat_response(resolved_chain.as_ref(), &mut outcome.response), @@ -3865,7 +3906,40 @@ where ), } }; - match ctx.chain.check_output(&synthesized).await { + let verdict = ctx.chain.check_output_non_segment(&synthesized).await; + let mut seg_counts = crate::redact::RedactionCounts::new(); + let verdict = crate::redact::moderate_body( + ctx.chain.as_ref(), + crate::redact::Direction::Output, + verdict, + &mut seg_counts, + |g| crate::redact::redact_chat_chunks(g, &mut pending), + ) + .await; + if !seg_counts.is_empty() { + // Bedrock masked the held chunks — rebuild the + // content-capture accumulator from the masked + // content channel (the sync redactor below can't + // reproduce a provider-side mask), keeping the + // original soft cap (#932 × AISIX-Cloud#947). + if let Some(cap) = content_cap { + let mut rebuilt = String::new(); + for c in pending.iter() { + if rebuilt.len() >= cap as usize { + break; + } + if let Some(t) = c.delta.content.as_deref() { + rebuilt.push_str(t); + } + } + guard.comp().response_text = rebuilt; + } + crate::redact::merge_counts( + &mut guard.comp().redacted_entity_counts, + seg_counts, + ); + } + match verdict { aisix_guardrails::GuardrailVerdict::Block { reason, guardrail_name } => { tracing::warn!( guardrail_hook = "output", diff --git a/crates/aisix-proxy/src/completions.rs b/crates/aisix-proxy/src/completions.rs index e377cd45..1b8ad16d 100644 --- a/crates/aisix-proxy/src/completions.rs +++ b/crates/aisix-proxy/src/completions.rs @@ -241,12 +241,26 @@ async fn dispatch( team_id: auth.key().team_id.as_deref(), }; let resolved_chain = state.guardrail_index.resolve(&guardrail_ctx); + let mut input_seg_counts = crate::redact::RedactionCounts::new(); if !resolved_chain.is_empty() { let chat = completions_input_to_chat(model_name, &body); + let verdict = + aisix_guardrails::Guardrail::check_input_non_segment(&resolved_chain, &chat).await; + // Segment pass: one Bedrock call over the prompt slots; an + // ANONYMIZE disposition writes the masked text back into the body + // (#932 bedrock follow-up). + let verdict = crate::redact::moderate_body( + &resolved_chain, + crate::redact::Direction::Input, + verdict, + &mut input_seg_counts, + |g| crate::redact::redact_completions_request(g, &mut body), + ) + .await; if let aisix_guardrails::GuardrailVerdict::Block { reason, guardrail_name, - } = aisix_guardrails::Guardrail::check_input(&resolved_chain, &chat).await + } = verdict { // Per #153 the matched-pattern detail stays in ops logs only. tracing::warn!( @@ -264,6 +278,7 @@ async fn dispatch( // #932: mask-action PII rules rewrite the prompt in place AFTER the // block check passes, BEFORE the body is forwarded upstream. let mut redactions = crate::redact::redact_completions_request(&resolved_chain, &mut body); + crate::redact::merge_counts(&mut redactions, input_seg_counts); // Content capture (AISIX-Cloud#947): the client-facing request body // (post-redaction, so masked PII stays masked in the exported content), @@ -327,6 +342,7 @@ async fn dispatch( // text into a synthetic ChatResponse and run the chain. The upstream // already billed (tokens committed above), so a block surfaces a // redacted 422 rather than the response. + let mut resp_json = resp_json; if !resolved_chain.is_empty() { let synth = ChatResponse { id: String::new(), @@ -335,10 +351,21 @@ async fn dispatch( finish_reason: FinishReason::Stop, usage: UsageStats::default(), }; + let verdict = + aisix_guardrails::Guardrail::check_output_non_segment(&resolved_chain, &synth) + .await; + let verdict = crate::redact::moderate_body( + &resolved_chain, + crate::redact::Direction::Output, + verdict, + &mut redactions, + |g| crate::redact::redact_completions_response(g, &mut resp_json), + ) + .await; if let aisix_guardrails::GuardrailVerdict::Block { reason, guardrail_name, - } = aisix_guardrails::Guardrail::check_output(&resolved_chain, &synth).await + } = verdict { // Per #153 the matched-pattern detail stays in ops logs only. tracing::warn!( @@ -376,7 +403,6 @@ async fn dispatch( // #932: mask-action PII rules rewrite the reply text AFTER the // block check passes. - let mut resp_json = resp_json; crate::redact::merge_counts( &mut redactions, crate::redact::redact_completions_response(&resolved_chain, &mut resp_json), diff --git a/crates/aisix-proxy/src/messages.rs b/crates/aisix-proxy/src/messages.rs index 11f33908..97d177e4 100644 --- a/crates/aisix-proxy/src/messages.rs +++ b/crates/aisix-proxy/src/messages.rs @@ -459,10 +459,26 @@ async fn dispatch( *applied_out = resolved_chain.applied().to_vec(); if !resolved_chain.is_empty() { if let Ok(chat) = aisix_provider_anthropic::parse_inbound_request(body) { + let verdict = aisix_guardrails::Guardrail::check_input_non_segment( + resolved_chain.as_ref(), + &chat, + ) + .await; + // Segment pass: one Bedrock call over the body's text slots; + // an ANONYMIZE disposition writes the masked text back into + // the Anthropic-native body (#932 bedrock follow-up). + let verdict = crate::redact::moderate_body( + resolved_chain.as_ref(), + crate::redact::Direction::Input, + verdict, + redactions_out, + |g| crate::redact::redact_anthropic_request(g, body), + ) + .await; if let aisix_guardrails::GuardrailVerdict::Block { reason, guardrail_name, - } = aisix_guardrails::Guardrail::check_input(resolved_chain.as_ref(), &chat).await + } = verdict { tracing::warn!( guardrail_hook = "input", @@ -1168,7 +1184,7 @@ async fn anthropic_passthrough_dispatch( // failures cool down the target — a body the bridge can't // parse is a real upstream problem worth taking out of // rotation, not a caller bug. - let json_body: Value = upstream_resp + let mut json_body: Value = upstream_resp .json() .await .map_err(|e| { @@ -1187,6 +1203,7 @@ async fn anthropic_passthrough_dispatch( // The body is forwarded verbatim, so extract its text (content // blocks + the raw content array, which covers tool_use args) into // a synthetic ChatResponse for inspection before returning it. + let mut output_seg_counts = crate::redact::RedactionCounts::new(); if !resolved_chain.is_empty() { if let Some(content) = json_body.get("content").and_then(|v| v.as_array()) { let mut out_text = String::new(); @@ -1210,11 +1227,23 @@ async fn anthropic_passthrough_dispatch( finish_reason: aisix_gateway::FinishReason::Stop, usage: aisix_gateway::UsageStats::new(0, 0), }; + let verdict = aisix_guardrails::Guardrail::check_output_non_segment( + resolved_chain.as_ref(), + &synth, + ) + .await; + let verdict = crate::redact::moderate_body( + resolved_chain.as_ref(), + crate::redact::Direction::Output, + verdict, + &mut output_seg_counts, + |g| crate::redact::redact_anthropic_response(g, &mut json_body), + ) + .await; if let aisix_guardrails::GuardrailVerdict::Block { reason, guardrail_name, - } = - aisix_guardrails::Guardrail::check_output(resolved_chain.as_ref(), &synth).await + } = verdict { tracing::warn!( guardrail_hook = "output", @@ -1233,7 +1262,6 @@ async fn anthropic_passthrough_dispatch( } // Restore the gateway-facing model name so callers see what they asked for. - let mut json_body = json_body; if let Some(m) = json_body.get_mut("model") { // If the upstream echoes the model name, rewrite to the gateway name. if m.as_str().map(|s| s == upstream_model).unwrap_or(false) { @@ -1243,8 +1271,9 @@ async fn anthropic_passthrough_dispatch( // #932: mask-action PII rules rewrite the passthrough response body // (text blocks + tool_use input) AFTER the block check passes. - let output_redactions = + let mut output_redactions = crate::redact::redact_anthropic_response(resolved_chain.as_ref(), &mut json_body); + crate::redact::merge_counts(&mut output_redactions, output_seg_counts); // Capture the prompt (the outbound request body) + assembled assistant // text for content-capturing exporters (gated). Built here, before @@ -1646,11 +1675,23 @@ async fn cross_provider_dispatch( // #448 (#22): run output guardrails on the cross-provider response // before rendering it back as Anthropic JSON — the response is // client-visible output just like /v1/chat/completions. + let mut output_seg_counts = crate::redact::RedactionCounts::new(); if !resolved_chain.is_empty() { + let verdict = + aisix_guardrails::Guardrail::check_output_non_segment(resolved_chain.as_ref(), &resp) + .await; + let verdict = crate::redact::moderate_body( + resolved_chain.as_ref(), + crate::redact::Direction::Output, + verdict, + &mut output_seg_counts, + |g| crate::redact::redact_chat_response(g, &mut resp), + ) + .await; if let aisix_guardrails::GuardrailVerdict::Block { reason, guardrail_name, - } = aisix_guardrails::Guardrail::check_output(resolved_chain.as_ref(), &resp).await + } = verdict { tracing::warn!( guardrail_hook = "output", @@ -1666,7 +1707,9 @@ async fn cross_provider_dispatch( // #932: mask-action PII rules rewrite the bridged response AFTER the // block check passes, BEFORE it is rendered back as Anthropic JSON. - let output_redactions = crate::redact::redact_chat_response(resolved_chain.as_ref(), &mut resp); + let mut output_redactions = + crate::redact::redact_chat_response(resolved_chain.as_ref(), &mut resp); + crate::redact::merge_counts(&mut output_redactions, output_seg_counts); let metrics = AnthropicUsageMetrics { prompt_tokens: resp.usage.prompt_tokens, @@ -1868,10 +1911,45 @@ fn build_anthropic_sse_stream( finish_reason: aisix_gateway::FinishReason::Stop, usage: aisix_gateway::UsageStats::new(0, 0), }; + let verdict = + aisix_guardrails::Guardrail::check_output_non_segment(chain.as_ref(), &synth) + .await; + let mut seg_counts = crate::redact::RedactionCounts::new(); + let verdict = crate::redact::moderate_body( + chain.as_ref(), + crate::redact::Direction::Output, + verdict, + &mut seg_counts, + |g| crate::redact::redact_chat_chunks(g, &mut held_chunks), + ) + .await; + if !seg_counts.is_empty() { + // Bedrock masked the held chunks — rebuild the content- + // capture accumulator from the masked content channel + // (the sync redactor below can't reproduce a provider- + // side mask), keeping the original soft cap + // (#932 × AISIX-Cloud#947). + if let Some(cap) = content_cap { + let mut rebuilt = String::new(); + for c in held_chunks.iter() { + if rebuilt.len() >= cap as usize { + break; + } + if let Some(t) = c.delta.content.as_deref() { + rebuilt.push_str(t); + } + } + guard.comp().response_text = rebuilt; + } + crate::redact::merge_counts( + &mut guard.comp().redacted_entity_counts, + seg_counts, + ); + } if let aisix_guardrails::GuardrailVerdict::Block { reason, guardrail_name, - } = aisix_guardrails::Guardrail::check_output(chain.as_ref(), &synth).await + } = verdict { tracing::warn!( guardrail_hook = "output", @@ -2631,10 +2709,51 @@ where finish_reason: aisix_gateway::FinishReason::Stop, usage: aisix_gateway::UsageStats::new(0, 0), }; + let verdict = + aisix_guardrails::Guardrail::check_output_non_segment(chain.as_ref(), &synth) + .await; + // Segment pass over the held SSE bytes. Only meaningful in + // hold-back mode (`held` is empty otherwise — and a chain + // with a segment member always folds to BufferFull, so a + // live-forward stream never carries one). + let mut seg_counts = crate::redact::RedactionCounts::new(); + let verdict = crate::redact::moderate_body( + chain.as_ref(), + crate::redact::Direction::Output, + verdict, + &mut seg_counts, + |g| match crate::redact::redact_anthropic_sse(g, &held) { + Some((rewritten, counts)) => { + held = rewritten; + counts + } + None => crate::redact::RedactionCounts::new(), + }, + ) + .await; + if !seg_counts.is_empty() { + // Bedrock masked the held bytes — rebuild the content- + // capture accumulator from the masked text channels + // (the sync redactor can't reproduce a provider-side + // mask) (#932 × AISIX-Cloud#947). + if let Some(cap) = content_cap { + let mut rebuilt = crate::redact::anthropic_sse_text(&held); + let mut cut = (cap as usize).min(rebuilt.len()); + while cut < rebuilt.len() && !rebuilt.is_char_boundary(cut) { + cut += 1; + } + rebuilt.truncate(cut); + guard.usage().response_text = rebuilt; + } + crate::redact::merge_counts( + &mut guard.usage().redacted_entity_counts, + seg_counts, + ); + } if let aisix_guardrails::GuardrailVerdict::Block { reason, guardrail_name, - } = aisix_guardrails::Guardrail::check_output(chain.as_ref(), &synth).await + } = verdict { tracing::warn!( guardrail_hook = "output", diff --git a/crates/aisix-proxy/src/redact.rs b/crates/aisix-proxy/src/redact.rs index 0cbcbd5d..6879cbb1 100644 --- a/crates/aisix-proxy/src/redact.rs +++ b/crates/aisix-proxy/src/redact.rs @@ -55,6 +55,182 @@ fn redact_str( } } +// ─── Remote segment moderation (kind=bedrock mask write-back) ──────────────── +// +// A Bedrock guardrail whose PII action is ANONYMIZE returns the masked +// replacement text from the SAME `ApplyGuardrail` call that yields the +// verdict — an async, whole-request rewrite that can't implement the sync +// per-field redact contract above. The bridge works in three walker +// passes over one wire body, all using the SAME wire-shape walker so slot +// enumeration order is identical by construction: +// +// 1. collect: a probe guardrail records every text slot the walker +// offers (rewriting nothing); +// 2. one remote call: the chain's segment fold sends the slots as one +// content block each and returns verdict + positionally-aligned +// masked texts; +// 3. apply: a second probe guardrail replaces slot i with masked[i]. +// +// Call sites pair this with `check_*_non_segment` so a segment-moderating +// member is consulted exactly once per hook. Families without a wire +// walker (embeddings, rerank, images, audio, passthrough, MCP) keep the +// plain `check_*` path, where an ANONYMIZE disposition still maps to +// Block — there is no write-back channel, and releasing the un-masked +// content would defeat the operator's policy. + +/// Marker count key [`SegmentApplier`] attaches to each rewritten slot. +/// Several walkers discard a rewrite whose counts are empty (their +/// "did anything change" gate); the marker makes those gates fire. It is +/// never surfaced: [`moderate_body`] discards the apply-walk's returned +/// counts and reports the provider's entity counts instead. +const SEGMENT_APPLY_MARKER: &str = "__segment_apply__"; + +/// Pass-1 probe: records every text slot the walker offers. Never +/// rewrites, so the body is bit-identical after the collect walk. +#[derive(Default)] +struct SegmentCollector { + texts: std::sync::Mutex>, +} + +impl SegmentCollector { + fn take(&self) -> Vec { + std::mem::take(&mut self.texts.lock().expect("collector poisoned")) + } + + fn record(&self, text: &str) -> Option { + self.texts + .lock() + .expect("collector poisoned") + .push(text.to_owned()); + None + } +} + +impl Guardrail for SegmentCollector { + fn name(&self) -> &'static str { + "segment-collector" + } + fn redacts_input(&self) -> bool { + true + } + fn redacts_output(&self) -> bool { + true + } + fn redact_input_text(&self, text: &str) -> Option { + self.record(text) + } + fn redact_output_text(&self, text: &str) -> Option { + self.record(text) + } +} + +/// Pass-3 probe: replaces the i-th offered slot with `masked[i]`. +/// Positional by construction — the walker offers slots in the same +/// order the collector recorded them (same walker, same body state). +struct SegmentApplier { + state: std::sync::Mutex, +} + +struct ApplierState { + masked: Vec, + cursor: usize, +} + +impl SegmentApplier { + fn new(masked: Vec) -> Self { + Self { + state: std::sync::Mutex::new(ApplierState { masked, cursor: 0 }), + } + } + + fn apply(&self, original: &str) -> Option { + let mut st = self.state.lock().expect("applier poisoned"); + let i = st.cursor; + st.cursor += 1; + match st.masked.get(i) { + Some(m) if m != original => Some(aisix_guardrails::Redaction { + text: m.clone(), + counts: std::iter::once((SEGMENT_APPLY_MARKER.to_owned(), 1)).collect(), + }), + _ => None, + } + } + + /// Warn when the apply walk offered a different slot count than the + /// mask carries — the extra/missing slots kept their originals (the + /// per-slot `get` above never misassigns), this is diagnostics only. + fn warn_if_misaligned(&self) { + let st = self.state.lock().expect("applier poisoned"); + if st.cursor != st.masked.len() { + tracing::warn!( + offered = st.cursor, + masked = st.masked.len(), + "segment apply walk drifted from collect walk; \ + unmatched slots kept their original text", + ); + } + } +} + +impl Guardrail for SegmentApplier { + fn name(&self) -> &'static str { + "segment-applier" + } + fn redacts_input(&self) -> bool { + true + } + fn redacts_output(&self) -> bool { + true + } + fn redact_input_text(&self, text: &str) -> Option { + self.apply(text) + } + fn redact_output_text(&self, text: &str) -> Option { + self.apply(text) + } +} + +/// Complete one hook's moderation over a wire body: fold the already-run +/// `check_*_non_segment` verdict with the remote segment pass. The +/// segment pass is skipped when the check already blocked (the request +/// is dead — don't burn a provider call) or when the chain has no +/// segment-moderating member (zero overhead for non-Bedrock chains). +/// Masked replacements are written back through `walk`; the provider's +/// entity counts merge into `counts_out` (they feed +/// `redacted_entity_counts`, names only — #932 no-leak). +pub async fn moderate_body( + chain: &dyn Guardrail, + dir: Direction, + non_segment_verdict: aisix_guardrails::GuardrailVerdict, + counts_out: &mut RedactionCounts, + mut walk: impl FnMut(&dyn Guardrail) -> RedactionCounts, +) -> aisix_guardrails::GuardrailVerdict { + if non_segment_verdict.is_block() || !chain.moderates_segments() { + return non_segment_verdict; + } + let collector = SegmentCollector::default(); + walk(&collector); + let texts = collector.take(); + if texts.is_empty() { + return non_segment_verdict; + } + let outcome = match dir { + Direction::Input => chain.moderate_input_segments(&texts).await, + Direction::Output => chain.moderate_output_segments(&texts).await, + }; + if !outcome.verdict.is_block() { + if let Some(masked) = outcome.masked { + let applier = SegmentApplier::new(masked); + // Marker counts are plumbing (see SEGMENT_APPLY_MARKER) — + // discard them; the provider counts below are the real ones. + let _ = walk(&applier); + applier.warn_if_misaligned(); + merge_counts(counts_out, outcome.counts); + } + } + non_segment_verdict.merged_with(outcome.verdict) +} + /// Rewrite one owned text field in place. No-op (and no allocation) when /// nothing matches. fn apply_to_string( @@ -863,6 +1039,76 @@ pub fn redact_anthropic_sse( Some((out, counts)) } +/// The concatenated TEXT-channel content of a buffered Anthropic-native +/// SSE stream (per content-block `index` order, `content_block_start` +/// head text included). Used to rebuild the content-capture accumulator +/// after a segment (provider-side) mask rewrote the held bytes — the +/// sync redactor can't reproduce a provider mask (#932 × AISIX-Cloud#947). +pub fn anthropic_sse_text(raw: &[u8]) -> String { + let (frames, _) = split_sse_frames(raw); + let mut channels: BTreeMap = BTreeMap::new(); + for frame in &frames { + let Some(data) = frame.data.as_ref() else { + continue; + }; + let index = data.get("index").and_then(Value::as_u64).unwrap_or(0); + let text = match data.get("type").and_then(Value::as_str) { + Some("content_block_delta") => data + .get("delta") + .filter(|d| d.get("type").and_then(Value::as_str) == Some("text_delta")) + .and_then(|d| d.get("text")) + .and_then(Value::as_str), + Some("content_block_start") => data + .get("content_block") + .and_then(|b| b.get("text")) + .and_then(Value::as_str), + _ => None, + }; + if let Some(t) = text { + channels.entry(index).or_default().push_str(t); + } + } + channels.into_values().collect() +} + +/// The concatenated `output_text` delta content of a buffered +/// `/v1/responses` SSE stream (channel order). Same capture-rebuild role +/// as [`anthropic_sse_text`]. +pub fn responses_sse_text(raw: &[u8]) -> String { + let (frames, _) = split_sse_frames(raw); + let mut channels: BTreeMap = BTreeMap::new(); + for frame in &frames { + let Some(data) = frame.data.as_ref() else { + continue; + }; + if data.get("type").and_then(Value::as_str) != Some("response.output_text.delta") { + continue; + } + let Some(t) = data.get("delta").and_then(Value::as_str) else { + continue; + }; + let key = match data.get("item_id").and_then(Value::as_str) { + Some(id) => format!( + "{id}/{}", + data.get("content_index") + .and_then(Value::as_u64) + .unwrap_or(0) + ), + None => format!( + "{}/{}", + data.get("output_index") + .and_then(Value::as_u64) + .unwrap_or(0), + data.get("content_index") + .and_then(Value::as_u64) + .unwrap_or(0) + ), + }; + channels.entry(key).or_default().push_str(t); + } + channels.into_values().collect() +} + // ─── Responses-API SSE rewrite ─────────────────────────────────────────────── /// Mask a fully-buffered Responses-API SSE byte stream (the `/v1/responses` diff --git a/crates/aisix-proxy/src/responses.rs b/crates/aisix-proxy/src/responses.rs index 3f367bd3..8bd3b06a 100644 --- a/crates/aisix-proxy/src/responses.rs +++ b/crates/aisix-proxy/src/responses.rs @@ -355,10 +355,24 @@ async fn dispatch( let resolved_chain = Arc::new(state.guardrail_index.resolve(&guardrail_ctx)); if !resolved_chain.is_empty() { let chat = responses_input_to_chat(&model_name, body); + let verdict = + aisix_guardrails::Guardrail::check_input_non_segment(resolved_chain.as_ref(), &chat) + .await; + // Segment pass: one Bedrock call over the body's text slots; an + // ANONYMIZE disposition writes the masked text back into the + // Responses body (#932 bedrock follow-up). + let verdict = crate::redact::moderate_body( + resolved_chain.as_ref(), + crate::redact::Direction::Input, + verdict, + redactions_out, + |g| crate::redact::redact_responses_request(g, body), + ) + .await; if let aisix_guardrails::GuardrailVerdict::Block { reason, guardrail_name, - } = aisix_guardrails::Guardrail::check_input(resolved_chain.as_ref(), &chat).await + } = verdict { // Per #153 the matched-pattern detail stays in ops logs only; the // wire envelope names only the guardrail that fired (#519 B.4b) @@ -951,10 +965,30 @@ async fn responses_to_target( } let out_text = responses_sse_output_text(&buf); let synth = synth_chat_response(&upstream_model, out_text); + let verdict = + aisix_guardrails::Guardrail::check_output_non_segment(chain, &synth).await; + // Segment pass over the held SSE frames: one Bedrock call; an + // ANONYMIZE disposition rewrites `buf` in place (#932 bedrock + // follow-up). The capture below reads the post-mask buffer. + let mut output_redactions = crate::redact::RedactionCounts::new(); + let verdict = crate::redact::moderate_body( + chain, + crate::redact::Direction::Output, + verdict, + &mut output_redactions, + |g| match crate::redact::redact_responses_sse(g, &buf) { + Some((rewritten, counts)) => { + buf = rewritten; + counts + } + None => crate::redact::RedactionCounts::new(), + }, + ) + .await; if let aisix_guardrails::GuardrailVerdict::Block { reason, guardrail_name, - } = aisix_guardrails::Guardrail::check_output(chain, &synth).await + } = verdict { // Per #153 the matched-pattern detail stays in ops logs only. tracing::warn!( @@ -969,10 +1003,9 @@ async fn responses_to_target( } // #932: the whole SSE response is held here — mask the frames // (channel reassembly) before anything reaches the wire. - let mut output_redactions = crate::redact::RedactionCounts::new(); let buf = match crate::redact::redact_responses_sse(chain, &buf) { Some((rewritten, counts)) => { - output_redactions = counts; + crate::redact::merge_counts(&mut output_redactions, counts); rewritten } None => buf, @@ -1167,12 +1200,24 @@ async fn responses_to_target( // configured output block isn't bypassable by calling /v1/responses // (the input half is enforced in `dispatch`). Only when an // output-hook guardrail is attached; otherwise this is a no-op. + let mut json_body = json_body; + let mut output_seg_counts = crate::redact::RedactionCounts::new(); if aisix_guardrails::Guardrail::runs_on_output(chain) { let synth = synth_chat_response(&upstream_model, responses_output_text(&json_body)); + let verdict = + aisix_guardrails::Guardrail::check_output_non_segment(chain, &synth).await; + let verdict = crate::redact::moderate_body( + chain, + crate::redact::Direction::Output, + verdict, + &mut output_seg_counts, + |g| crate::redact::redact_responses_response(g, &mut json_body), + ) + .await; if let aisix_guardrails::GuardrailVerdict::Block { reason, guardrail_name, - } = aisix_guardrails::Guardrail::check_output(chain, &synth).await + } = verdict { // Per #153 the matched-pattern detail stays in ops logs only. tracing::warn!( @@ -1209,8 +1254,8 @@ async fn responses_to_target( // #932: mask-action PII rules rewrite the response body AFTER the // block check passes. - let mut json_body = json_body; - let output_redactions = crate::redact::redact_responses_response(chain, &mut json_body); + let mut output_redactions = crate::redact::redact_responses_response(chain, &mut json_body); + crate::redact::merge_counts(&mut output_redactions, output_seg_counts); // Content capture (AISIX-Cloud#947): the assistant's assembled output // text, read from the POST-redaction body so masked PII stays masked @@ -1528,11 +1573,22 @@ async fn responses_cross_provider_to_target( // #719: run output guardrails on the bridged response before re-encoding // it as Responses JSON — the assistant text + tool calls are // client-visible output, scanned the same way /v1/chat/completions does. + let mut output_seg_counts = crate::redact::RedactionCounts::new(); if aisix_guardrails::Guardrail::runs_on_output(chain.as_ref()) { + let verdict = + aisix_guardrails::Guardrail::check_output_non_segment(chain.as_ref(), &resp).await; + let verdict = crate::redact::moderate_body( + chain.as_ref(), + crate::redact::Direction::Output, + verdict, + &mut output_seg_counts, + |g| crate::redact::redact_chat_response(g, &mut resp), + ) + .await; if let aisix_guardrails::GuardrailVerdict::Block { reason, guardrail_name, - } = aisix_guardrails::Guardrail::check_output(chain.as_ref(), &resp).await + } = verdict { tracing::warn!( guardrail_hook = "output", @@ -1566,7 +1622,8 @@ async fn responses_cross_provider_to_target( // #932: mask-action PII rules rewrite the bridged response AFTER the // block check passes, BEFORE it is re-encoded as Responses JSON. - let output_redactions = crate::redact::redact_chat_response(chain.as_ref(), &mut resp); + let mut output_redactions = crate::redact::redact_chat_response(chain.as_ref(), &mut resp); + crate::redact::merge_counts(&mut output_redactions, output_seg_counts); let created_at = chrono::Utc::now().timestamp(); let json_body = crate::responses_bridge::chat_response_to_responses_json( diff --git a/crates/aisix-proxy/src/responses_bridge.rs b/crates/aisix-proxy/src/responses_bridge.rs index 78bedd79..d963fd9b 100644 --- a/crates/aisix-proxy/src/responses_bridge.rs +++ b/crates/aisix-proxy/src/responses_bridge.rs @@ -1082,9 +1082,52 @@ pub fn build_responses_bridge_stream( finish_reason: FinishReason::Stop, usage: UsageStats::new(0, 0), }; - if let aisix_guardrails::GuardrailVerdict::Block { reason, guardrail_name } = - aisix_guardrails::Guardrail::check_output(chain.as_ref(), &synth).await - { + let verdict = + aisix_guardrails::Guardrail::check_output_non_segment(chain.as_ref(), &synth) + .await; + // Segment pass over the held SSE frames: one Bedrock call; an + // ANONYMIZE disposition rewrites the held bytes (#932 bedrock + // follow-up). + let mut seg_counts = crate::redact::RedactionCounts::new(); + let mut joined: Vec = Vec::with_capacity(held_bytes); + for b in &held { + joined.extend_from_slice(b); + } + let mut seg_rewrote = false; + let verdict = crate::redact::moderate_body( + chain.as_ref(), + crate::redact::Direction::Output, + verdict, + &mut seg_counts, + |g| match crate::redact::redact_responses_sse(g, &joined) { + Some((rewritten, counts)) => { + joined = rewritten; + seg_rewrote = true; + counts + } + None => crate::redact::RedactionCounts::new(), + }, + ) + .await; + if !seg_counts.is_empty() { + // Bedrock masked the held bytes — rebuild the content- + // capture accumulator from the masked text channels, + // keeping the original soft cap (#932 × AISIX-Cloud#947). + if let Some(cap) = content_cap { + let mut rebuilt = crate::redact::responses_sse_text(&joined); + let mut cut = (cap as usize).min(rebuilt.len()); + while cut < rebuilt.len() && !rebuilt.is_char_boundary(cut) { + cut += 1; + } + rebuilt.truncate(cut); + guard.comp().response_text = rebuilt; + } + crate::redact::merge_counts( + &mut guard.comp().redacted_entity_counts, + seg_counts, + ); + } + if let aisix_guardrails::GuardrailVerdict::Block { reason, guardrail_name } = verdict { tracing::warn!( guardrail_hook = "output", model = %model_label, @@ -1095,6 +1138,9 @@ pub fn build_responses_bridge_stream( yield Ok(bytes::Bytes::from(guardrail_error_frame(guardrail_name.as_deref()))); return; } + if seg_rewrote { + held = vec![bytes::Bytes::from(joined)]; + } } // Passed (#932): mask the held SSE frames (channel reassembly) // before release, then hand them to the client. From 07d4e907d6b7f5cd58203b6d72afbb8c64488144 Mon Sep 17 00:00:00 2001 From: Jarvis Date: Mon, 6 Jul 2026 12:36:03 +0800 Subject: [PATCH 3/4] test(bedrock): unit + e2e coverage for ANONYMIZE segment write-back MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - redact.rs unit tests: collect→call→apply positional round trip over the chat walker (flat content / text block / tool-arg JSON slots each index-stamped), Block leaves the body untouched, prior-Block and no-segment-member skip the remote call, Anthropic-SSE channel masking via the marker-count gate, and the SSE capture-rebuild helpers. - e2e (vitest, mock Bedrock ApplyGuardrail): request masked per slot before the upstream (multi-block INPUT call pinned), reply masked before the caller (non-streaming + split-across-chunks streaming), BLOCKED entity still 422s before the upstream, misaligned outputs[] fall back to originals-and-continue, /v1/messages Anthropic-native body masked before the bridge. --- crates/aisix-proxy/src/redact.rs | 242 ++++++++++ .../cases/bedrock-anonymize-mask-e2e.test.ts | 429 ++++++++++++++++++ 2 files changed, 671 insertions(+) create mode 100644 tests/e2e/src/cases/bedrock-anonymize-mask-e2e.test.ts diff --git a/crates/aisix-proxy/src/redact.rs b/crates/aisix-proxy/src/redact.rs index 6879cbb1..ded0b28a 100644 --- a/crates/aisix-proxy/src/redact.rs +++ b/crates/aisix-proxy/src/redact.rs @@ -1786,4 +1786,246 @@ mod tests { assert_eq!(counts.get("email"), Some(&1)); assert!(redact_transcription_response(chain.as_ref(), b"all clean\n").is_none()); } + + // ── remote segment moderation (#932 bedrock follow-up) ────────────── + + use aisix_guardrails::{GuardrailVerdict, SegmentsOutcome}; + + /// Stub of a Bedrock-style segment moderator: masks slot i to + /// `""` — index-stamped so a positional mix-up is + /// unmissable — and reports a fixed entity count. `verdict` lets the + /// block/bypass paths be exercised; `panic_if_called` pins the + /// skip-when-already-blocked contract. + struct StubSegments { + verdict: GuardrailVerdict, + mask: bool, + panic_if_called: bool, + } + + impl StubSegments { + fn masker() -> Self { + Self { + verdict: GuardrailVerdict::Allow, + mask: true, + panic_if_called: false, + } + } + } + + #[async_trait::async_trait] + impl Guardrail for StubSegments { + fn name(&self) -> &'static str { + "stub-segments" + } + fn moderates_segments(&self) -> bool { + true + } + async fn moderate_input_segments(&self, texts: &[String]) -> SegmentsOutcome { + self.moderate(texts) + } + async fn moderate_output_segments(&self, texts: &[String]) -> SegmentsOutcome { + self.moderate(texts) + } + } + + impl StubSegments { + fn moderate(&self, texts: &[String]) -> SegmentsOutcome { + if self.panic_if_called { + panic!("segment moderator must not be called on this path"); + } + let mut counts = RedactionCounts::new(); + counts.insert("EMAIL".to_owned(), 1); + SegmentsOutcome { + verdict: self.verdict.clone(), + masked: self.mask.then(|| { + texts + .iter() + .enumerate() + .map(|(i, t)| format!("", t.to_uppercase())) + .collect() + }), + counts, + } + } + } + + fn seg_chain(stub: StubSegments) -> GuardrailChain { + GuardrailChain::new(vec![Arc::new(stub)]) + } + + /// The collect→call→apply round trip over the chat walker: every slot + /// kind (flat content, text block, tool-call JSON argument) gets its + /// OWN positionally-matched mask, and the provider counts — not the + /// applier's plumbing marker — land in `counts_out`. + #[tokio::test] + async fn moderate_body_masks_chat_slots_positionally() { + let chain = seg_chain(StubSegments::masker()); + let mut req: ChatFormat = serde_json::from_value(json!({ + "model": "m", + "messages": [ + {"role": "user", "content": "first slot"}, + {"role": "user", "content": "", "content_blocks": [ + {"type": "text", "text": "second slot"} + ]}, + {"role": "assistant", "content": null, "tool_calls": [ + {"index": 0, "function": {"name": "send", "arguments": "{\"to\":\"third slot\"}"}} + ]} + ] + })) + .unwrap(); + let mut counts = RedactionCounts::new(); + let verdict = moderate_body( + &chain, + Direction::Input, + GuardrailVerdict::Allow, + &mut counts, + |g| redact_chat_format(g, &mut req), + ) + .await; + assert_eq!(verdict, GuardrailVerdict::Allow); + assert_eq!( + req.messages[0].content.as_deref(), + Some(""), + "flat content = slot 0", + ); + assert_eq!( + req.messages[1].content_blocks.as_ref().unwrap()[0]["text"], + "", + "text block = slot 1", + ); + assert_eq!( + req.messages[2].extra["tool_calls"][0]["function"]["arguments"] + .as_str() + .unwrap(), + "{\"to\":\"\"}", + "tool-arg inner string = slot 2 (marker counts must fire the \ + json-encoded rewrite gate)", + ); + assert_eq!(counts.get("EMAIL"), Some(&1), "provider counts merged"); + assert!( + !counts.keys().any(|k| k.starts_with("__")), + "the applier's plumbing marker must never leak into telemetry counts", + ); + } + + /// A Block from the segment pass leaves the body untouched (no mask + /// write-back on a dead request) and propagates the verdict. + #[tokio::test] + async fn moderate_body_block_leaves_body_untouched() { + let chain = seg_chain(StubSegments { + verdict: GuardrailVerdict::block("pii blocked"), + mask: true, + panic_if_called: false, + }); + let mut req: ChatFormat = serde_json::from_value(json!({ + "model": "m", + "messages": [{"role": "user", "content": "original"}] + })) + .unwrap(); + let mut counts = RedactionCounts::new(); + let verdict = moderate_body( + &chain, + Direction::Input, + GuardrailVerdict::Allow, + &mut counts, + |g| redact_chat_format(g, &mut req), + ) + .await; + assert!(verdict.is_block()); + assert_eq!(req.messages[0].content.as_deref(), Some("original")); + assert!(counts.is_empty(), "no counts on a blocked request"); + } + + /// An already-blocked prior verdict skips the remote call entirely + /// (the request is dead — don't burn a provider call), and a chain + /// with no segment member is a no-op. + #[tokio::test] + async fn moderate_body_skips_remote_when_blocked_or_absent() { + let chain = seg_chain(StubSegments { + verdict: GuardrailVerdict::Allow, + mask: false, + panic_if_called: true, + }); + let mut req: ChatFormat = serde_json::from_value(json!({ + "model": "m", + "messages": [{"role": "user", "content": "x"}] + })) + .unwrap(); + let mut counts = RedactionCounts::new(); + let verdict = moderate_body( + &chain, + Direction::Input, + GuardrailVerdict::block("already blocked"), + &mut counts, + |g| redact_chat_format(g, &mut req), + ) + .await; + assert!(verdict.is_block(), "prior Block passes through"); + + // A sync-only (non-segment) chain never enters the pass. + let sync_only = both(); + let verdict = moderate_body( + sync_only.as_ref(), + Direction::Input, + GuardrailVerdict::Allow, + &mut counts, + |_| panic!("walk must not run when no segment member exists"), + ) + .await; + assert_eq!(verdict, GuardrailVerdict::Allow); + } + + /// The round trip through the Anthropic SSE walker: the masked + /// channel text lands on the channel's first frame (later frames + /// empty) even though the applier only returns marker counts — the + /// gate that discards count-less SSE rewrites must fire on them. + #[tokio::test] + async fn moderate_body_masks_anthropic_sse_channels() { + let chain = seg_chain(StubSegments::masker()); + let mut held: Vec = concat!( + "event: content_block_delta\n", + "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"hello \"}}\n\n", + "event: content_block_delta\n", + "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"world\"}}\n\n", + ) + .as_bytes() + .to_vec(); + let mut counts = RedactionCounts::new(); + let verdict = moderate_body( + &chain, + Direction::Output, + GuardrailVerdict::Allow, + &mut counts, + |g| match redact_anthropic_sse(g, &held) { + Some((rewritten, c)) => { + held = rewritten; + c + } + None => RedactionCounts::new(), + }, + ) + .await; + assert_eq!(verdict, GuardrailVerdict::Allow); + let out = String::from_utf8(held.clone()).unwrap(); + assert!( + out.contains(""), + "channel text masked as one positional slot: {out}", + ); + assert_eq!(counts.get("EMAIL"), Some(&1)); + // The capture-rebuild helper reads the masked channel back. + assert_eq!(anthropic_sse_text(&held), ""); + } + + /// `responses_sse_text` assembles `output_text` deltas per channel — + /// the capture-rebuild source after a segment mask. + #[test] + fn responses_sse_text_assembles_channels() { + let raw = concat!( + "event: response.output_text.delta\n", + "data: {\"type\":\"response.output_text.delta\",\"item_id\":\"i1\",\"content_index\":0,\"delta\":\"foo \"}\n\n", + "event: response.output_text.delta\n", + "data: {\"type\":\"response.output_text.delta\",\"item_id\":\"i1\",\"content_index\":0,\"delta\":\"bar\"}\n\n", + ); + assert_eq!(responses_sse_text(raw.as_bytes()), "foo bar"); + } } diff --git a/tests/e2e/src/cases/bedrock-anonymize-mask-e2e.test.ts b/tests/e2e/src/cases/bedrock-anonymize-mask-e2e.test.ts new file mode 100644 index 00000000..704286bd --- /dev/null +++ b/tests/e2e/src/cases/bedrock-anonymize-mask-e2e.test.ts @@ -0,0 +1,429 @@ +import { createHash } from "node:crypto"; +import { createServer, type Server } from "node:http"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + AdminClient, + EtcdClient, + pickFreePort, + spawnApp, + startOpenAiUpstream, + waitConfigPropagation, + type OpenAiUpstream, + type SpawnedApp, +} from "../harness/index.js"; + +// E2E: Bedrock guardrail ANONYMIZE write-back (#932 bedrock follow-up). +// +// Bedrock reports `action=GUARDRAIL_INTERVENED` for BOTH a hard block and +// a PII anonymization; the per-entity `assessments[]` actions tell them +// apart. Pre-fix the DP blocked on ANY intervention. Now, on the +// chat-shaped families, an ANONYMIZE disposition masks-and-continues: +// +// - INPUT: the request's text slots go up as one content block each in a +// single ApplyGuardrail call, and Bedrock's `outputs[i]` replaces slot i +// before the request reaches the upstream (verified via the mock +// upstream's received body). +// - OUTPUT (non-streaming + streaming hold-back): the model's reply is +// masked before it reaches the caller. +// - Hard block (BLOCKED entity) still blocks with the standard 422 / +// error-frame envelope. +// - Defensive fallback (LiteLLM `_merge_masked_texts` semantics): masked +// outputs that don't align 1:1 with the input blocks are NOT applied — +// originals pass through, the request continues. + +const CALLER = "sk-bedrock-mask-caller"; +const STREAM_CALLER = "sk-bedrock-mask-stream-caller"; +const hash = (s: string) => createHash("sha256").update(s).digest("hex"); + +const EMAIL = "alice@example.com"; +const EMAIL_RE = /[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}/g; + +interface BedrockCall { + source: string; + texts: string[]; +} + +interface MockBedrock { + url: string; + calls: BedrockCall[]; + close(): Promise; +} + +// Content-driven ApplyGuardrail mock: +// - any block containing "BLOCKME" → INTERVENED, BLOCKED pii entity +// - any block containing "MISALIGN" → INTERVENED, ANONYMIZED, but ONE +// merged output for N input blocks (deliberately misaligned) +// - any block containing an email → INTERVENED, ANONYMIZED, one output +// per input block with emails replaced by {EMAIL} (the real contract) +// - otherwise → NONE +async function startMockBedrock(): Promise { + const calls: BedrockCall[] = []; + const server: Server = createServer((req, res) => { + let raw = ""; + req.on("data", (c: Buffer) => (raw += c.toString("utf8"))); + req.on("end", () => { + let source = ""; + let texts: string[] = []; + try { + const body = JSON.parse(raw) as { + source?: string; + content?: Array<{ text?: { text?: string } }>; + }; + source = body.source ?? ""; + texts = (body.content ?? []).map((c) => c.text?.text ?? ""); + } catch { + // fall through with empty texts — answered as NONE below + } + calls.push({ source, texts }); + + const anonymizedAssessment = { + sensitiveInformationPolicy: { + piiEntities: [{ match: EMAIL, type: "EMAIL", action: "ANONYMIZED" }], + regexes: [], + }, + }; + let payload: unknown; + if (texts.some((t) => t.includes("BLOCKME"))) { + payload = { + action: "GUARDRAIL_INTERVENED", + outputs: [], + assessments: [ + { + sensitiveInformationPolicy: { + piiEntities: [ + { match: "BLOCKME", type: "NAME", action: "BLOCKED" }, + ], + regexes: [], + }, + }, + ], + }; + } else if (texts.some((t) => t.includes("MISALIGN"))) { + payload = { + action: "GUARDRAIL_INTERVENED", + outputs: [{ text: texts.join(" ").replace(EMAIL_RE, "{EMAIL}") }], + assessments: [anonymizedAssessment], + }; + } else if (texts.some((t) => t.includes("@"))) { + payload = { + action: "GUARDRAIL_INTERVENED", + outputs: texts.map((t) => ({ text: t.replace(EMAIL_RE, "{EMAIL}") })), + assessments: [anonymizedAssessment], + }; + } else { + payload = { action: "NONE", outputs: [] }; + } + res.statusCode = 200; + res.setHeader("content-type", "application/json"); + res.end(JSON.stringify(payload)); + }); + }); + const port = await pickFreePort(); + await new Promise((resolve) => + server.listen(port, "127.0.0.1", resolve), + ); + return { + url: `http://127.0.0.1:${port}`, + calls, + async close() { + await new Promise((resolve, reject) => { + server.close((err) => (err ? reject(err) : resolve())); + }); + }, + }; +} + +describe("bedrock guardrail ANONYMIZE write-back (#932 follow-up)", () => { + let etcdReachable = false; + let upstream: OpenAiUpstream | undefined; + let streamUpstream: OpenAiUpstream | undefined; + let bedrock: MockBedrock | undefined; + let app: SpawnedApp | undefined; + let admin: AdminClient | undefined; + + beforeAll(async () => { + etcdReachable = await new EtcdClient().ping(); + if (!etcdReachable) return; + + // Non-streaming upstream: the canned reply CONTAINS an email so the + // output mask has something to rewrite. + upstream = await startOpenAiUpstream({ + nonStreamBody: { + id: "cmpl-bmask", + object: "chat.completion", + created: 1, + model: "gpt-4o-mini", + choices: [ + { + index: 0, + message: { + role: "assistant", + content: `you can reach the customer at ${EMAIL} today`, + }, + finish_reason: "stop", + }, + ], + usage: { prompt_tokens: 5, completion_tokens: 8, total_tokens: 13 }, + }, + }); + + // Streaming upstream: the email split across two delta chunks — only + // the hold-back channel reassembly (one OUTPUT scan over the joined + // text) can catch the span. + streamUpstream = await startOpenAiUpstream({ + streamEvents: [ + '{"id":"strm-bmask","object":"chat.completion.chunk","model":"gpt-4o-mini","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]}', + '{"id":"strm-bmask","object":"chat.completion.chunk","model":"gpt-4o-mini","choices":[{"index":0,"delta":{"content":"mail alice@exam"},"finish_reason":null}]}', + '{"id":"strm-bmask","object":"chat.completion.chunk","model":"gpt-4o-mini","choices":[{"index":0,"delta":{"content":"ple.com now"},"finish_reason":null}]}', + '{"id":"strm-bmask","object":"chat.completion.chunk","model":"gpt-4o-mini","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}', + "[DONE]", + ], + eventDelayMs: 10, + }); + + bedrock = await startMockBedrock(); + app = await spawnApp({ extra: { bedrock_endpoint_url: bedrock.url } }); + admin = new AdminClient(app.adminUrl, app.adminKey); + + const pk = await admin.createProviderKey({ + display_name: "bmask-pk", + secret: "sk-mock", + api_base: `${upstream.baseUrl}/v1`, + }); + await admin.createModel({ + display_name: "bmask-e2e", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: pk.id, + }); + await admin.createApiKey({ + key_hash: hash(CALLER), + allowed_models: ["bmask-e2e"], + }); + + const streamPk = await admin.createProviderKey({ + display_name: "bmask-stream-pk", + secret: "sk-mock", + api_base: `${streamUpstream.baseUrl}/v1`, + }); + await admin.createModel({ + display_name: "bmask-stream-e2e", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: streamPk.id, + }); + await admin.createApiKey({ + key_hash: hash(STREAM_CALLER), + allowed_models: ["bmask-stream-e2e"], + }); + + await admin.json("POST", "/admin/v1/guardrails", { + name: "gr-bedrock-mask", + enabled: true, + hook_point: "both", + fail_open: false, + kind: "bedrock", + guardrail_id: "bmaskgr00001", + guardrail_version: "DRAFT", + region: "us-east-1", + aws_credentials: { + kind: "static", + access_key_id: "AKIDBEDROCKMASK00001", + secret_access_key: "secret-bedrock-mask", + }, + latency_mode: { kind: "serial" }, + }); + + // Ready once a clean chat passes AND the guardrail fired on INPUT. + await waitConfigPropagation(async () => { + try { + const r = await chat(CALLER, "bmask-e2e", [ + { role: "user", content: "warmup all clean" }, + ]); + await r.text(); + return ( + r.status === 200 && + bedrock!.calls.some((c) => c.source === "INPUT") + ); + } catch { + return false; + } + }); + }); + + afterAll(async () => { + await app?.exit(); + await bedrock?.close(); + await upstream?.close(); + await streamUpstream?.close(); + }); + + function chat( + caller: string, + model: string, + messages: Array<{ role: string; content: string }>, + stream = false, + ) { + return fetch(`${app!.proxyUrl}/v1/chat/completions`, { + method: "POST", + headers: { + authorization: `Bearer ${caller}`, + "content-type": "application/json", + }, + body: JSON.stringify({ model, messages, ...(stream ? { stream: true } : {}) }), + }); + } + + test( + "ANONYMIZE masks the request per slot and the reply before the caller", + async (ctx) => { + if (!etcdReachable) { + ctx.skip(); + return; + } + const upstreamBefore = upstream!.receivedRequests.length; + const r = await chat(CALLER, "bmask-e2e", [ + { role: "system", content: "be helpful" }, + { role: "user", content: `contact ${EMAIL} please` }, + ]); + expect(r.status).toBe(200); + const body = (await r.json()) as { + choices: Array<{ message: { content: string } }>; + }; + + // INPUT went up as one content block per message (Plan D), not a + // joined blob. + const inputCall = bedrock!.calls + .filter((c) => c.source === "INPUT") + .at(-1)!; + expect(inputCall.texts).toEqual([ + "be helpful", + `contact ${EMAIL} please`, + ]); + + // The upstream saw the MASKED prompt — the raw email never left. + const sent = upstream!.receivedRequests.slice(upstreamBefore).at(-1)!; + expect(sent.body).toContain("contact {EMAIL} please"); + expect(sent.body).not.toContain(EMAIL); + + // The reply's email is masked before the caller sees it. + expect(body.choices[0].message.content).toBe( + "you can reach the customer at {EMAIL} today", + ); + expect( + bedrock!.calls.some((c) => c.source === "OUTPUT"), + "output hook must have scanned the reply", + ).toBe(true); + }, + 60_000, + ); + + test( + "a BLOCKED entity still blocks before the upstream", + async (ctx) => { + if (!etcdReachable) { + ctx.skip(); + return; + } + const upstreamBefore = upstream!.receivedRequests.length; + const r = await chat(CALLER, "bmask-e2e", [ + { role: "user", content: "please BLOCKME now" }, + ]); + expect(r.status).toBe(422); + const body = await r.text(); + expect(body).toContain("content policy"); + expect(body).not.toContain("BLOCKME"); + expect(upstream!.receivedRequests.length).toBe(upstreamBefore); + }, + 60_000, + ); + + test( + "misaligned masked outputs are not applied — originals pass through", + async (ctx) => { + if (!etcdReachable) { + ctx.skip(); + return; + } + const upstreamBefore = upstream!.receivedRequests.length; + // Two slots, but the mock answers with ONE merged output — the + // fallback must keep the originals and continue. + const r = await chat(CALLER, "bmask-e2e", [ + { role: "user", content: `MISALIGN write to bob@example.org` }, + { role: "user", content: "second slot" }, + ]); + expect(r.status).toBe(200); + await r.text(); + const sent = upstream!.receivedRequests.slice(upstreamBefore).at(-1)!; + expect(sent.body).toContain("bob@example.org"); + expect(sent.body).not.toContain("{EMAIL}"); + }, + 60_000, + ); + + test( + "streaming: the reply is masked across chunk boundaries via hold-back", + async (ctx) => { + if (!etcdReachable) { + ctx.skip(); + return; + } + const r = await chat( + STREAM_CALLER, + "bmask-stream-e2e", + [{ role: "user", content: "stream something" }], + true, + ); + expect(r.status).toBe(200); + const raw = await r.text(); + const content = raw + .split("\n") + .filter((l) => l.startsWith("data: ") && !l.includes("[DONE]")) + .map((l) => { + try { + const j = JSON.parse(l.slice("data: ".length)) as { + choices?: Array<{ delta?: { content?: string } }>; + }; + return j.choices?.[0]?.delta?.content ?? ""; + } catch { + return ""; + } + }) + .join(""); + expect(content).toBe("mail {EMAIL} now"); + expect(content).not.toContain("alice@"); + }, + 60_000, + ); + + test( + "/v1/messages: the Anthropic-native body is masked before the bridge", + async (ctx) => { + if (!etcdReachable) { + ctx.skip(); + return; + } + const upstreamBefore = upstream!.receivedRequests.length; + const r = await fetch(`${app!.proxyUrl}/v1/messages`, { + method: "POST", + headers: { + authorization: `Bearer ${CALLER}`, + "content-type": "application/json", + }, + body: JSON.stringify({ + model: "bmask-e2e", + max_tokens: 64, + messages: [ + { role: "user", content: `write to carol@example.dev today` }, + ], + }), + }); + expect(r.status).toBe(200); + await r.text(); + const sent = upstream!.receivedRequests.slice(upstreamBefore).at(-1)!; + expect(sent.body).toContain("write to {EMAIL} today"); + expect(sent.body).not.toContain("carol@example.dev"); + }, + 60_000, + ); +}); From cc2f2bca958338665d21ce5eed72f3d07b87ee00 Mon Sep 17 00:00:00 2001 From: Jarvis Date: Mon, 6 Jul 2026 13:34:48 +0800 Subject: [PATCH 4/4] fix(bedrock): gate hard-block on per-entry action; review fixes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CodeRabbit findings on #719: - topic/content/word policy entries also have a detect-only mode (action=NONE, observability metadata): gate the hard-block decision on action=BLOCKED per entry instead of entry presence, matching LiteLLM's per-entry check — a detect-mode entry no longer turns an ANONYMIZE mask into a block. - chain fold: merge a member's entity counts only when its mask is accepted — a refused (misaligned) mask must not inflate redacted_entity_counts. - responses_sse_text: concatenate channels in first-seen emission order, not item-id lexicographic order. --- crates/aisix-guardrails/src/bedrock.rs | 67 ++++++++++++++++++++++---- crates/aisix-guardrails/src/chain.rs | 14 +++++- crates/aisix-proxy/src/redact.rs | 27 +++++++++-- 3 files changed, 93 insertions(+), 15 deletions(-) diff --git a/crates/aisix-guardrails/src/bedrock.rs b/crates/aisix-guardrails/src/bedrock.rs index ef558e54..679ed88a 100644 --- a/crates/aisix-guardrails/src/bedrock.rs +++ b/crates/aisix-guardrails/src/bedrock.rs @@ -392,15 +392,34 @@ fn anonymized_counts(resp: &ApplyGuardrailOutput) -> std::collections::BTreeMap< } /// True if the assessment carries any BLOCKING disposition — a -/// topic/content/word policy hit (these have no anonymize mode, so their -/// presence is a block), a contextual-grounding filter that BLOCKED, or a -/// PII/regex entity whose action is BLOCKED (as opposed to ANONYMIZED). +/// topic/content/word/contextual-grounding entry whose `action` is +/// BLOCKED, or a PII/regex entity whose action is BLOCKED (as opposed to +/// ANONYMIZED). Every policy family also has a detect-only mode +/// (`action = NONE`): entries are returned for observability but nothing +/// was suppressed, so their mere presence must NOT be read as a block — +/// mirrors LiteLLM's per-entry BLOCKED check. fn assessment_has_hard_block(a: &GuardrailAssessment) -> bool { - let topic_blocked = a.topic_policy().is_some_and(|p| !p.topics().is_empty()); - let content_blocked = a.content_policy().is_some_and(|p| !p.filters().is_empty()); - let word_blocked = a - .word_policy() - .is_some_and(|p| !p.custom_words().is_empty() || !p.managed_word_lists().is_empty()); + use aws_sdk_bedrockruntime::types::{ + GuardrailContentPolicyAction, GuardrailTopicPolicyAction, GuardrailWordPolicyAction, + }; + let topic_blocked = a.topic_policy().is_some_and(|p| { + p.topics() + .iter() + .any(|t| *t.action() == GuardrailTopicPolicyAction::Blocked) + }); + let content_blocked = a.content_policy().is_some_and(|p| { + p.filters() + .iter() + .any(|f| *f.action() == GuardrailContentPolicyAction::Blocked) + }); + let word_blocked = a.word_policy().is_some_and(|p| { + p.custom_words() + .iter() + .any(|w| *w.action() == GuardrailWordPolicyAction::Blocked) + || p.managed_word_lists() + .iter() + .any(|w| *w.action() == GuardrailWordPolicyAction::Blocked) + }); let grounding_blocked = a.contextual_grounding_policy().is_some_and(|p| { p.filters().iter().any(|f| { *f.action() @@ -632,11 +651,11 @@ mod tests { .build() } - fn topic() -> GuardrailAssessment { + fn topic_with(action: GuardrailTopicPolicyAction) -> GuardrailAssessment { let t = GuardrailTopic::builder() .name("blocked-topic") .r#type(GuardrailTopicType::Deny) - .action(GuardrailTopicPolicyAction::Blocked) + .action(action) .build() .unwrap(); let tp = GuardrailTopicPolicyAssessment::builder() @@ -646,6 +665,10 @@ mod tests { GuardrailAssessment::builder().topic_policy(tp).build() } + fn topic() -> GuardrailAssessment { + topic_with(GuardrailTopicPolicyAction::Blocked) + } + #[test] fn action_none_is_allow() { let r = resp(GuardrailAction::None, vec![], vec![]); @@ -708,6 +731,30 @@ mod tests { assert!(!assessment_has_hard_block(&pii(PiiAction::Anonymized))); assert!(assessment_has_hard_block(&pii(PiiAction::Blocked))); assert!(assessment_has_hard_block(&topic())); + // Detect-only (`action=NONE`) entries are observability + // metadata, not a block. + assert!(!assessment_has_hard_block(&topic_with( + GuardrailTopicPolicyAction::None + ))); + } + + /// A detect-mode (`action=NONE`) topic entry alongside an + /// ANONYMIZED PII entity must classify as Mask — the topic + /// policy observed but did not suppress anything. + #[test] + fn detect_only_topic_does_not_turn_mask_into_block() { + let r = resp( + GuardrailAction::GuardrailIntervened, + vec!["contact {EMAIL}"], + vec![ + topic_with(GuardrailTopicPolicyAction::None), + pii(PiiAction::Anonymized), + ], + ); + assert_eq!( + classify_response(&r, "gid"), + BedrockOutcome::Mask(vec!["contact {EMAIL}".to_owned()]) + ); } /// Positional integrity: an empty `outputs[]` entry is preserved, diff --git a/crates/aisix-guardrails/src/chain.rs b/crates/aisix-guardrails/src/chain.rs index cf46bb0b..e6b2bebd 100644 --- a/crates/aisix-guardrails/src/chain.rs +++ b/crates/aisix-guardrails/src/chain.rs @@ -343,8 +343,12 @@ async fn fold_segments(members: &[ChainMember], texts: &[String], input: bool) - if let Some(new_masked) = outcome.masked { // Implementations uphold alignment with THEIR input; refuse a // drifted length here so a broken member can't desync slots. + // Counts merge ONLY with an accepted mask — they describe + // APPLIED anonymization (`redacted_entity_counts`), so a + // refused mask must not inflate them. if new_masked.len() == src.len() { masked = Some(new_masked); + Redaction::merge_counts(&mut counts, &outcome.counts); } else { tracing::warn!( member = %m.name, @@ -354,7 +358,6 @@ async fn fold_segments(members: &[ChainMember], texts: &[String], input: bool) - ); } } - Redaction::merge_counts(&mut counts, &outcome.counts); } SegmentsOutcome { verdict: match bypass { @@ -780,10 +783,12 @@ mod tests { true } async fn moderate_input_segments(&self, _texts: &[String]) -> crate::SegmentsOutcome { + let mut counts = std::collections::BTreeMap::new(); + counts.insert("EMAIL".to_owned(), 3); crate::SegmentsOutcome { verdict: GuardrailVerdict::Allow, masked: Some(vec!["only-one".to_owned()]), - counts: std::collections::BTreeMap::new(), + counts, } } } @@ -791,6 +796,11 @@ mod tests { let texts = vec!["a".to_owned(), "b".to_owned()]; let out = chain.moderate_input_segments(&texts).await; assert_eq!(out.masked, None, "drifted mask must be refused"); + assert!( + out.counts.is_empty(), + "a refused mask's counts describe anonymization that was NOT \ + applied — they must not reach redacted_entity_counts", + ); assert_eq!(out.verdict, GuardrailVerdict::Allow); } diff --git a/crates/aisix-proxy/src/redact.rs b/crates/aisix-proxy/src/redact.rs index ded0b28a..5f3720b7 100644 --- a/crates/aisix-proxy/src/redact.rs +++ b/crates/aisix-proxy/src/redact.rs @@ -1076,7 +1076,9 @@ pub fn anthropic_sse_text(raw: &[u8]) -> String { /// as [`anthropic_sse_text`]. pub fn responses_sse_text(raw: &[u8]) -> String { let (frames, _) = split_sse_frames(raw); - let mut channels: BTreeMap = BTreeMap::new(); + // First-seen channel order (NOT key order): the rebuilt capture must + // read in the order the client saw the channels emitted. + let mut channels: Vec<(String, String)> = Vec::new(); for frame in &frames { let Some(data) = frame.data.as_ref() else { continue; @@ -1104,9 +1106,12 @@ pub fn responses_sse_text(raw: &[u8]) -> String { .unwrap_or(0) ), }; - channels.entry(key).or_default().push_str(t); + match channels.iter_mut().find(|(k, _)| *k == key) { + Some((_, buf)) => buf.push_str(t), + None => channels.push((key, t.to_owned())), + } } - channels.into_values().collect() + channels.into_iter().map(|(_, text)| text).collect() } // ─── Responses-API SSE rewrite ─────────────────────────────────────────────── @@ -2028,4 +2033,20 @@ mod tests { ); assert_eq!(responses_sse_text(raw.as_bytes()), "foo bar"); } + + /// Channels concatenate in first-seen (emission) order, not item-id + /// lexicographic order — the rebuilt capture must read like the + /// stream the client saw. + #[test] + fn responses_sse_text_preserves_emission_order() { + let raw = concat!( + "event: response.output_text.delta\n", + "data: {\"type\":\"response.output_text.delta\",\"item_id\":\"zzz\",\"content_index\":0,\"delta\":\"first \"}\n\n", + "event: response.output_text.delta\n", + "data: {\"type\":\"response.output_text.delta\",\"item_id\":\"aaa\",\"content_index\":0,\"delta\":\"second\"}\n\n", + "event: response.output_text.delta\n", + "data: {\"type\":\"response.output_text.delta\",\"item_id\":\"zzz\",\"content_index\":0,\"delta\":\"more\"}\n\n", + ); + assert_eq!(responses_sse_text(raw.as_bytes()), "first moresecond"); + } }