fix(proxy): run output guardrails on streaming responses - #222
Merged
Conversation
## Problem
Pre-fix the streaming path (`build_sse_stream` in
`crates/aisix-proxy/src/chat.rs`) skipped output guardrails entirely.
A `kind: "keyword"` deny-list configured on `hook_point: "output"`
fired correctly on non-streaming responses but was trivially
bypassable by setting `stream: true` — the gateway forwarded SSE
chunks verbatim and never accumulated the assistant text for
guardrail evaluation. So a forbidden literal in a streamed
completion reached the caller untouched.
A guardrail that only fires when someone forgets to enable a flag
isn't a security control. The reference implementation in this
category treats streaming output guardrails as a baseline
requirement (every streaming chat-completion request goes through
the guardrail dispatcher).
## Fix
Buffer-then-check at end-of-stream:
1. `build_sse_stream` now takes an optional `StreamGuardrailContext`
(a typed wrapper around `Arc<dyn Guardrail>` + the model name for
tracing). When set, the stream loop accumulates `chunk.delta.content`
into a `String` buffer.
2. After the loop completes cleanly (no upstream error), if a
guardrail context is present, synthesize a `ChatResponse` from the
accumulated content + the StreamCompletion's tracked usage stats
and call `chain.check_output(&resp).await`.
3. On `Block { reason }`:
- Mirror #153's wire-level redaction: emit an SSE `event: error`
frame with the OpenAI envelope shape (`error.type:
"content_filter"`, generic `"response blocked by content policy"`
message — no matched-literal leak).
- Set `errored = true` so the terminal `[DONE]` is suppressed
(per docs §5 abnormal-termination contract).
- Emit `tracing::warn!` with the rich verdict reason (operator-
side log carries the matched detail; wire-side does not).
- Set `comp.guardrail_blocked = true` so the post-stream
telemetry callback records the block on `usage_events`.
The call site in `dispatch_streaming` always passes `Some(ctx)` —
the chain itself short-circuits when no policies are configured, so
the per-chunk overhead for guardrail-free deployments is just a
single `Option::is_some()` check + an unused `String` allocation.
## Trade-off (documented)
Per-chunk evaluation is wrong for blocking guardrails: by the time
chunk N matches the forbidden literal, chunks 1..N-1 have already
been emitted — the secret leaks regardless. Buffer-then-check
trades latency-to-first-completion for the security guarantee.
Streaming masking-style guardrails (where partial leakage is
repaired by rewriting tokens) could use a different cadence; that's
a v2 concern.
## Tests
**Rust unit (`crates/aisix-proxy/src/lib.rs`)** — new
`streaming_output_guardrail_blocks_with_sse_error_event_and_no_done`
test that drains raw response bytes through the full proxy stack
and asserts:
- `!wire.contains("data: [DONE]")` — terminal sentinel suppressed
- `wire.contains("event: error")` — error frame emitted
- error-frame data parses as JSON with `error.type === "content_filter"`
- error message equals the redacted static string
- error message does NOT contain the matched literal
**E2E (`tests/e2e/src/cases/guardrail-output-e2e.test.ts`)** — added
streaming case to the existing output-guardrail describe block.
Sets up a separate streaming upstream + Model + caller key sharing
the existing env-wide guardrail policy. Sends `stream: true` request
via raw `fetch`, asserts the wire-level shape (no `[DONE]`,
`event: error` present, OpenAI envelope, redacted message, no
forbidden literal in the error envelope, upstream IS hit per
buffer-then-check semantics).
## Verification
- `cargo test -p aisix-proxy --lib`: **158/158 passing** (incl. new test)
- `cargo clippy --workspace --lib --tests -- -D warnings`: clean
- `cargo fmt --check`: clean
- `pnpm tsc --noEmit` (e2e): clean
## References
- Issue: #204 (`bug` + `vulnerability`)
- `docs/api-proxy.md` §5 (abnormal-termination contract — same SSE
shape this fix uses for guardrail blocks)
- #153 (non-streaming guardrail redaction — same wire-level
contract mirrored to streaming)
- #198 (`error_frame_payload` helper — reused here)
ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Free Run ID: 📒 Files selected for processing (4)
✅ Files skipped from review due to trivial changes (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughThe pull request adds end-of-stream output guardrail evaluation for streaming chat responses. A new ChangesStreaming Output Guardrail Blocking
🎯 4 (Complex) | ⏱️ ~45 minutes Note 🎁 Summarized by CodeRabbit FreeYour organization is on the Free plan. CodeRabbit will generate a high-level summary and a walkthrough for each pull request. For a comprehensive line-by-line review, please upgrade your subscription to CodeRabbit Pro by visiting https://app.coderabbit.ai/login. Comment |
There was a problem hiding this comment.
Pull request overview
This PR closes #204 by ensuring output guardrails are enforced for streaming /v1/chat/completions responses (previously bypassable with stream: true). It implements a buffer-then-check strategy: accumulate streamed assistant content, run check_output after the upstream stream completes, and on Block emit an SSE event: error with a redacted OpenAI-style error envelope and omit [DONE].
Changes:
- Run output guardrails at end-of-stream for streaming chat completions; emit SSE
event: errorand suppress[DONE]on block. - Add
guardrail_blockedpropagation into streaming telemetry (UsageEvent.guardrail_blocked). - Add Rust unit + E2E coverage pinning the wire-level contract for streaming blocks (error event + no
[DONE]+ redacted message/type).
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
crates/aisix-proxy/src/chat.rs |
Adds streaming output-guardrail buffering/checking and propagates guardrail_blocked into stream completion telemetry. |
crates/aisix-proxy/src/lib.rs |
Adds a unit test asserting streaming guardrail blocks emit SSE error + omit [DONE] and redact matched literal in the error envelope. |
tests/e2e/src/cases/guardrail-output-e2e.test.ts |
Extends existing guardrail E2E to cover the streaming wire shape end-to-end. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
512
to
519
| // Per #204: pass the gateway's guardrail chain so the | ||
| // streaming path can run output guardrails at end-of-stream | ||
| // (buffer-then-check). Mirrors the non-streaming | ||
| // `state.guardrails.check_output(...)` call site. | ||
| let stream_guardrail = Some(StreamGuardrailContext { | ||
| chain: Arc::clone(&state.guardrails), | ||
| model_name: req.model.clone(), | ||
| }); |
Comment on lines
+1210
to
+1219
| // Per #204: accumulate the assistant's content across chunks | ||
| // so the output guardrail can evaluate the full response at | ||
| // end-of-stream. Allocates only when an output guardrail is | ||
| // configured AND the upstream actually emits content; for | ||
| // requests without an output guardrail this is a no-op | ||
| // borrow of the empty string. | ||
| let mut content_buffer = if output_guardrail.is_some() { | ||
| Some(String::new()) | ||
| } else { | ||
| None |
Comment on lines
+128
to
+134
| // Add the streaming Model alias to the existing caller's allow-list. | ||
| await admin.createApiKey({ | ||
| key_hash: createHash("sha256") | ||
| .update(`${CALLER_PLAINTEXT}-stream`) | ||
| .digest("hex"), | ||
| allowed_models: ["gr-out-stream-e2e"], | ||
| }); |
…ath skip, Allow e2e Audit fixes on PR #222: - HIGH-1: Output `Bypass` verdict at end-of-stream was silently dropped from telemetry. Pre-fix the streaming arm only matched `Block`; `Bypass { reason }` fell through to `[DONE]` with no side effect, while the non-streaming path captures Bypass into `usage_events.bypass_reason` so operators can audit which policy fail-opened. Added `bypass_reason: String` field to `StreamCompletion`, populated on the Bypass arm with first-bypass- wins semantics. The on_complete closure merges with the input- side `bypass_reason_for_telem` snapshot before emitting telemetry. - MEDIUM-2: No fast-path skip when the guardrail chain is empty — `String` allocated per chunk on every streaming request even for guardrail-free deployments. Added `is_empty(&self) -> bool` default-method to the `Guardrail` trait (returns `false`, preserving safe behavior for custom impls); `GuardrailChain` overrides to return `self.guardrails.is_empty()`. The streaming call site in `chat.rs::dispatch_streaming` now passes `Some(StreamGuardrailContext)` only when the chain is non-empty; guardrail-free deployments skip both the per-chunk content accumulation and the post-loop synthesized-ChatResponse build. - MEDIUM-3: E2E did not cover the streaming-Allow path — a regression that ALWAYS blocks streaming would have passed the existing block-case test. Added a companion case in `guardrail-output-e2e.test.ts` with a separate clean upstream + Model + caller that emits content NOT containing `FORBIDDEN_WORD`. Asserts the wire shape: terminal `[DONE]` IS present, no `event: error`, full assistant content reaches the caller. - MEDIUM-1 (client-disconnect bypasses guardrail check) → filed as #223. The post-loop check sits after `while let Some(item) = upstream.next().await`; client disconnect drops the generator at the last yield suspension, so the post-loop code never executes and `comp.guardrail_blocked` stays `false` even when the buffered content would have blocked. Telemetry under-counts disconnects-as-blocks. Out of scope for this PR (already a step forward vs pre-fix "every streaming request bypasses the guardrail"); recommended fix shape is to spawn a detached task from the `CompleteOnDrop` Drop impl. - LOW-1, LOW-2 (cosmetic — model_name field, message-content clone): reviewed, deferred. Both are micro-optimizations on a path already gated behind the is_empty() fast-path, so they don't affect the dominant guardrail-free deployment. Verification: - cargo test -p aisix-proxy --lib: 158/158 passing - cargo test -p aisix-guardrails --lib: 42/42 passing - cargo clippy --workspace --lib --tests -- -D warnings: clean - cargo fmt --check: clean - pnpm tsc --noEmit (e2e): clean
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
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.
Closes #204.
Problem
Pre-fix the streaming path skipped output guardrails entirely. A
kind: "keyword"deny-list configured onhook_point: "output"fired correctly on non-streaming responses but was trivially bypassable by settingstream: true— the gateway forwarded SSE chunks verbatim and never accumulated the assistant text for guardrail evaluation. A forbidden literal in a streamed completion reached the caller untouched.A security control that only fires when someone forgets to enable a flag isn't a security control. The de-facto reference implementation in this category treats streaming output guardrails as a baseline requirement.
Fix — buffer-then-check at end-of-stream
build_sse_streamnow takes an optionalStreamGuardrailContextand accumulateschunk.delta.contentinto aStringbuffer. After the upstream stream completes cleanly (no upstream error), the gateway callschain.check_output(...)on a synthesizedChatResponse. OnBlock:event: errorframe with OpenAI envelope (error.type: "content_filter", redacted"response blocked by content policy"message — mirrors bug: output guardrail error message echoes the matched forbidden literal back to caller #153's non-streaming wire contract)[DONE]so SDK consumers detect the truncation (per docs §5)tracing::warn!with the rich verdict reason (operator-side carries matched detail; wire-side does not)comp.guardrail_blocked = trueso post-stream telemetry records the blockWhy buffer-then-check vs per-chunk
Per-chunk evaluation is wrong for blocking guardrails: by the time chunk N matches the forbidden literal, chunks 1..N-1 have already been emitted — the secret leaks regardless. Buffer-then-check trades latency-to-first-completion for the security guarantee.
Streaming masking-style guardrails (where partial leakage is repaired by rewriting tokens, e.g. PII redaction) could use a different cadence; that's a v2 concern, not in scope here. The
kind: "keyword"deny-list use case in #204 is the blocking variant.Trade-off — partial leakage of the chunks BEFORE the block fires
The pre-emitted
data: ...chunks (chunks 1..N) DO carry the partial assistant content to the caller before the buffer-then-check completes. The security guarantee is "the stream did NOT close cleanly, AND an SSEevent: errorframe surfaces the block" — not "byte-perfect prevention of all leakage". Preventing every single byte from reaching the wire would require holding ALL chunks server-side until the check fires, which negates streaming's latency-to-first-token benefit.This trade-off matches the de-facto reference implementation's approach for buffer-then-check guardrails.
Tests
Rust unit (
crates/aisix-proxy/src/lib.rs) —streaming_output_guardrail_blocks_with_sse_error_event_and_no_done:!wire.contains("data: [DONE]")— terminal sentinel suppressedwire.contains("event: error")— error frame emittederror.type === "content_filter"E2E (
tests/e2e/src/cases/guardrail-output-e2e.test.ts) — added streaming case to the existing describe. Separate streaming upstream + Model + caller key sharing the env-wide guardrail policy. Rawfetchwithstream: true. Asserts the wire-level shape end-to-end.Verification
cargo test -p aisix-proxy --lib: 158/158 passingcargo test -p aisix-guardrails --lib: 42/42 passingcargo clippy --workspace --lib --tests -- -D warnings: cleancargo fmt --check: cleanpnpm tsc --noEmit(e2e): cleanAudit follow-ups (commit b799761)
Independent audit on initial push surfaced 1 HIGH + 3 MEDIUM + 2 LOW. Resolved as follows:
HIGH-1 → fixed: output
Bypassverdict at end-of-stream was silently dropped from telemetry. Addedbypass_reason: Stringfield toStreamCompletion; on Bypass arm, capture with first-bypass-wins semantics (matches non-streaming convention). The on_complete closure merges with the input-sidebypass_reason_for_telemsnapshot before emitting telemetry.MEDIUM-2 → fixed: no fast-path skip when the guardrail chain is empty. Added
is_empty(&self) -> booldefault-method toGuardrailtrait (returnsfalse— safe default for custom impls);GuardrailChainoverrides. Streaming call site now passesSome(StreamGuardrailContext)only when!state.guardrails.is_empty(); guardrail-free deployments skip both per-chunk content accumulation and the post-loop synthesized-ChatResponsebuild.MEDIUM-3 → fixed: e2e didn't cover the streaming-Allow path. Added a companion case in
guardrail-output-e2e.test.tswith a separate clean upstream + Model + caller emitting content NOT containingFORBIDDEN_WORD. Asserts terminal[DONE]IS present, noevent: error, full assistant content reaches the wire. A regression that ALWAYS blocks would have passed the existing block-case test alone — this companion catches it.MEDIUM-1 (client-disconnect bypasses guardrail check) → filed as bug: streaming output guardrail check is skipped on client disconnect (telemetry under-counts blocks) #223: the post-loop check sits after
while let Some(item) = upstream.next().await; client disconnect drops the generator at the last yield suspension, so the post-loop code never runs andcomp.guardrail_blockedstaysfalseeven when the buffered content would have blocked. Telemetry under-counts disconnect-as-block. Out of scope for this PR (already a step forward vs the pre-fix "every streaming request bypasses the guardrail"). The recommended fix shape is to spawn a detached task from theCompleteOnDropDrop impl — tracked as a separate change because it requires async-from-Drop scaffolding that isn't in this PR.LOW-1 (cosmetic —
model_namefield used only in tracing line) → deferred: current shape is fine.LOW-2 (cosmetic —
ChatMessage::assistant(content.clone())clones) → deferred: gated behind the is_empty() fast-path now (M2 fix), so the clone only happens on guardrail-active deployments where the cost is negligible vs upstream I/O.References
bug+vulnerability)docs/api-proxy.md§5 (abnormal-termination contract — reused for guardrail block shape)error_frame_payloadhelper — reused here)Summary by CodeRabbit
Bug Fixes
Tests