feat(ai): guard streamed output in the Vercel AI SDK middleware (wrapStream) - #5
Merged
Merged
Conversation
aegisMiddleware() implemented transformParams (input) and wrapGenerate (non-streaming output) but not wrapStream. streamText() — the most common production path — routes through wrapStream, so every output guard (LLM02 PII, LLM05 improper-output, LLM06 disclosure, LLM08 agency) was silently bypassed for streamed responses, even though the middleware was attached and its docstring claimed it scanned generated text. Add wrapStream: it pipes the provider's stream through a TransformStream that feeds text deltas into the existing sliding-window stream guard (createStreamGuard), erroring the stream with AegisBlockedError the moment a detector trips — before the offending delta reaches the consumer. This mirrors the fail-closed semantics of wrapGenerate and catches violations that straddle chunk boundaries. Adds tests: benign streamed output passes through unchanged; PII that completes across deltas aborts the stream. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
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.
Why
aegisMiddleware()implementedtransformParams(input scan) andwrapGenerate(non-streaming output scan) but notwrapStream. In the Vercel AI SDK,streamText()— the most common production path — routes through the middleware'swrapStreamhook, notwrapGenerate. So when a model streams, every output guard (LLM02 PII, LLM05 improper-output, LLM06 disclosure, LLM08 agency) was silently bypassed even though the middleware was attached, and the file's docstring's claim that it scans "the generated text after" was false for streaming.For a security middleware, silently not guarding the most-used path is a serious gap.
What
wrapStream: pipe the provider'sstreamthrough aTransformStreamthat feedstext-deltachunks into the existing sliding-window stream guard (createStreamGuard), andcontroller.error(AegisBlockedError)the moment a detector trips — before the offending delta is emitted. Non-text parts pass through untouched.guardTextStream, so it catches violations that straddle chunk boundaries (e.g.admin@+example.com), and mirrorswrapGenerate's fail-closed throw-on-violation semantics.generateTextandstreamTextare covered.Verification
Two new tests (benign streamed output passes through unchanged; PII completing across deltas aborts the stream). Full suite green:
typecheck·lint·test(310 pass) ·build·bench(thresholds met, 0% FP) ·perf(p95 within budget).🤖 Generated with Claude Code