feat(bedrock): honor ANONYMIZE — mask-and-continue instead of block - #719
Merged
Conversation
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
…tion
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).
- 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.
|
Warning Review limit reachedYou’ve reached a temporary PR review limit under our Fair Usage Limits Policy. Next review available in: 2 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (3)
📝 WalkthroughWalkthroughThis PR adds masking-aware segment moderation for Bedrock guardrails, introducing ChangesSegment-aware anonymization pipeline
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Dispatch as Proxy Dispatch
participant Chain as GuardrailChain
participant Bedrock as BedrockGuardrail
participant ModerateBody as redact::moderate_body
Dispatch->>Chain: check_input_non_segment(req)
Chain-->>Dispatch: verdict (Allow/Block/Bypass)
Dispatch->>ModerateBody: moderate_body(chain, dir, verdict, walk)
ModerateBody->>Chain: collect text slots
ModerateBody->>Chain: moderate_input_segments(texts)
Chain->>Bedrock: apply_segments(texts)
Bedrock-->>Chain: SegmentsOutcome(masked, counts)
Chain-->>ModerateBody: merged SegmentsOutcome
ModerateBody->>Dispatch: rewritten body + merged counts
Dispatch->>Dispatch: proceed with masked request/response
Possibly related PRs
Important Pre-merge checks failedPlease resolve all errors before merging. Addressing warnings is optional. ❌ Failed checks (1 error, 1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (1)
tests/e2e/src/cases/bedrock-anonymize-mask-e2e.test.ts (1)
294-303: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
bedrock!.callslookup relies on hidden test-order assumptions.
bedrock.callsis a single array populated acrossbeforeAll's warmup call and every subsequent test, never reset or scoped. Grabbing.filter(c => c.source === "INPUT").at(-1)only returns this test's call because tests happen to run sequentially in file order. Elsewhere in this same file (Line 305, 357, 423) you correctly snapshotreceivedRequests.lengthbefore issuing the request and slice from that baseline — the same pattern isn't applied here, making this assertion fragile to reordering or an inserted guardrail retry.♻️ Suggested fix: snapshot before the call, like the upstream-request assertions
+ const bedrockCallsBefore = bedrock!.calls.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 inputCall = bedrock!.calls + .slice(bedrockCallsBefore) .filter((c) => c.source === "INPUT") .at(-1)!;As per path instructions for
**/*.{test,spec,e2e}.{js,ts,jsx,tsx}: "Avoid explicit dependencies between tests and hidden execution order assumptions."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/e2e/src/cases/bedrock-anonymize-mask-e2e.test.ts` around lines 294 - 303, The INPUT assertion is depending on the shared bedrock.calls array and hidden test execution order, which makes it fragile. In the affected test, snapshot the current bedrock.calls length before issuing the request, then assert only against the new calls added by this test instead of using .filter(...).at(-1). Use the existing request-baseline pattern already used elsewhere in bedrock-anonymize-mask-e2e.test.ts so the check is scoped to this test case and not prior warmup or retries.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/aisix-guardrails/src/bedrock.rs`:
- Around line 194-201: The Bedrock request building in
moderate_*_segments()/send() is still adding GuardrailContentBlock::Text("") for
empty segment slots, which can break the whole request. Update the text loop in
bedrock.rs to skip or specially handle empty strings before calling
GuardrailTextBlock::builder() and req.content(...), while keeping outputs[i]
aligned with the original segment positions.
- Around line 398-420: The hard-block check in assessment_has_hard_block is
treating any topic/content/word match as a block even when the policy action is
NONE. Update the topic_policy, content_policy, and word_policy branches to gate
on each policy’s action the same way the contextual_grounding_policy and
sensitive_information_policy branches do, so detect-only assessments don’t force
classify_response to Block. Add or adjust a test covering NONE mode to verify
mask/continue flows remain non-blocking.
In `@crates/aisix-guardrails/src/chain.rs`:
- Around line 311-368: The fold_segments function is merging Redaction counts
even when a member’s masked output is rejected for being misaligned, which can
overstate applied anonymization. Update fold_segments so Redaction::merge_counts
only runs when the new masked output is accepted (or otherwise explicitly
handled), using the existing masked/new_masked logic and GuardrailVerdict flow
to keep counts aligned with applied masks. Add a test with a drifting stub that
returns non-empty counts and a misaligned mask to verify redacted_entity_counts
is not incremented unless the mask is actually accepted.
In `@crates/aisix-proxy/src/redact.rs`:
- Around line 1090-1106: The emission-order bug is in the `responses_sse_text`
aggregation path that builds the `key` from `item_id`, `output_index`, and
`content_index` before iterating via `BTreeMap::into_values()`. Replace the
lexicographic map-order reconstruction with an order-preserving approach, such
as tracking first-seen insertion order or explicitly sorting by `output_index`
and `content_index` before rebuilding `comp.response_text`, so multi-item
Responses streams reflect the original emission order.
---
Nitpick comments:
In `@tests/e2e/src/cases/bedrock-anonymize-mask-e2e.test.ts`:
- Around line 294-303: The INPUT assertion is depending on the shared
bedrock.calls array and hidden test execution order, which makes it fragile. In
the affected test, snapshot the current bedrock.calls length before issuing the
request, then assert only against the new calls added by this test instead of
using .filter(...).at(-1). Use the existing request-baseline pattern already
used elsewhere in bedrock-anonymize-mask-e2e.test.ts so the check is scoped to
this test case and not prior warmup or retries.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 3867497c-0de9-43ac-945f-efba34951f7a
📒 Files selected for processing (10)
crates/aisix-guardrails/src/bedrock.rscrates/aisix-guardrails/src/chain.rscrates/aisix-guardrails/src/lib.rscrates/aisix-proxy/src/chat.rscrates/aisix-proxy/src/completions.rscrates/aisix-proxy/src/messages.rscrates/aisix-proxy/src/redact.rscrates/aisix-proxy/src/responses.rscrates/aisix-proxy/src/responses_bridge.rstests/e2e/src/cases/bedrock-anonymize-mask-e2e.test.ts
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
A Bedrock guardrail configured with PII action Anonymize is unusable through the gateway: Bedrock reports
action=GUARDRAIL_INTERVENEDfor both a hard block and an anonymization, andkind=bedrocktreated every intervention as a block. Operators who chose "mask the PII and continue" got a 422 instead of the masked conversation.Implementation
One
ApplyGuardrailcall per hook now carries the request's/response's text slots as one content block each (instead of a single joined blob), so Bedrock'soutputs[i]aligns positionally with sloti— the same contract LiteLLM relies on. The response is classified via the per-entityassessments[]actions:BLOCKEDdisposition (topic/content/word/contextual-grounding policy hit, or aBLOCKEDPII/regex entity) → still blocks;ANONYMIZED-only → the maskedoutputs[]texts are written back into the wire body and the request/response continues;outputs[](count ≠ input blocks) → defensive fallback: the mask is not applied, originals stand, the request continues (LiteLLM_merge_masked_textssemantics — never misapply masked text to the wrong slot).Mechanics: the
Guardrailtrait gainsmoderate_{input,output}_segments(verdict + positional mask + entity counts from one remote call) andcheck_{input,output}_non_segment; call sites pair the two so a segment member is consulted exactly once per hook. Slot enumeration reuses the #932 wire walkers via collect/apply probe guardrails (SegmentCollector/SegmentApplierinredact.rs), so collect and write-back order are identical by construction across every wire shape, including streaming hold-back (chunk channels and SSE frames). Masked entity types (e.g.EMAIL) merge intoredacted_entity_counts; assessmentmatchvalues are never read (#153/#932 no-leak).Scope: the chat-shaped families —
/v1/chat/completions(incl. cache hits, ensemble, streaming),/v1/messages(passthrough + bridge, non-streaming + SSE),/v1/responses(passthrough + bridge, non-streaming + SSE),/v1/completions. Families with no mask write-back channel (embeddings, rerank, images, audio, raw passthrough, MCP) keep blob mode, where ANONYMIZE still maps to Block — releasing un-masked content there would defeat the operator's policy.Behavior change
Bedrock guardrails whose matched policy action is ANONYMIZE previously blocked the request/response; they now mask and continue on the families above. Hard-block policies, fail-open/closed handling, timeouts, and throttle tagging are unchanged. No new config surface (masking is driven by the Bedrock guardrail's own configuration, matching LiteLLM), hence no CP changes.
Streaming content capture (AISIX-Cloud#947 invariant) is preserved: after a provider-side mask, the capture accumulator is rebuilt from the masked chunks/SSE channels so exporters never see pre-mask text.
Tests
apply_segmentswiremock tests (one call with N blocks, positional mask, misalign fallback, hard block, fail-open/closed, blob mode still blocks on anonymize), chain segment folds (compose, attribute, refuse drifted masks).bedrock-anonymize-mask-e2e.test.ts, mock Bedrock): request masked per slot before the upstream, reply masked non-streaming + across streaming chunk boundaries, BLOCKED still 422s pre-upstream, misalign fallback passes originals,/v1/messagesbody masked before the bridge.Follow-up to the
kind=piiwork in AISIX-Cloud#932 (Bedrock external one-way masking half).🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Bug Fixes