diff --git a/src/adapters/ai.ts b/src/adapters/ai.ts index 860a41c..3766697 100644 --- a/src/adapters/ai.ts +++ b/src/adapters/ai.ts @@ -3,7 +3,8 @@ * * `aegisMiddleware()` is a `LanguageModelMiddleware` you pass to `wrapLanguageModel`. * It scans the prompt (LLM01/LLM10) before the model runs and the generated text - * (LLM02/LLM06/LLM08) after, throwing `AegisBlockedError` on a violation. + * (LLM02/LLM06/LLM08) after — for BOTH `generateText` (`wrapGenerate`) and + * `streamText` (`wrapStream`) — throwing `AegisBlockedError` on a violation. * * Zero runtime dependency on the AI SDK — the import is type-only. * @@ -14,6 +15,7 @@ */ import type { LanguageModelMiddleware } from "ai"; import { createAegisGuard } from "../aegis-guard.js"; +import { createStreamGuard } from "../stream.js"; import type { AegisOptions } from "../types.js"; import { AegisBlockedError } from "../errors.js"; @@ -101,5 +103,33 @@ export function aegisMiddleware( } return result; }, + wrapStream: async ({ doStream }) => { + // streamText() routes through wrapStream — NOT wrapGenerate — so without + // this hook every output guard (PII/disclosure/improper-output/agency) was + // silently bypassed for streamed responses. Feed accumulated text deltas + // through the same sliding-window stream guard and error the stream the + // moment a detector trips, before the offending delta reaches the consumer. + const result = await doStream(); + const sg = createStreamGuard({ enabled: true, ...options, scope: "output" }); + const stream = result.stream.pipeThrough( + new TransformStream({ + async transform(part, controller) { + if ( + isRecord(part) && + part["type"] === "text-delta" && + typeof part["delta"] === "string" + ) { + const r = await sg.push(part["delta"]); + if (r.blocked) { + controller.error(new AegisBlockedError("output", r.result)); + return; + } + } + controller.enqueue(part); + }, + }), + ); + return { ...result, stream }; + }, }; } diff --git a/tests/adapters.test.ts b/tests/adapters.test.ts index eae6b64..7ed6854 100644 --- a/tests/adapters.test.ts +++ b/tests/adapters.test.ts @@ -226,4 +226,51 @@ describe("aegisMiddleware (Vercel AI SDK)", () => { }); expect(out).toBe(result); }); + + // streamText() routes through wrapStream, not wrapGenerate; these guard against + // the regression where streamed output was never scanned. + const delta = (text: string): any => ({ type: "text-delta", id: "1", delta: text }); + function streamOf(parts: any[]): any { + return { + stream: new ReadableStream({ + start(controller) { + for (const p of parts) controller.enqueue(p); + controller.close(); + }, + }), + }; + } + async function drain(stream: ReadableStream): Promise { + const reader = stream.getReader(); + let out = ""; + for (;;) { + const { done, value } = await reader.read(); + if (done) break; + if (value && value.type === "text-delta") out += value.delta; + } + return out; + } + + it("passes benign streamed output through unchanged via wrapStream", async () => { + const mw = aegisMiddleware(); + const out = await mw.wrapStream!({ + doStream: async () => streamOf([delta("All "), delta("systems "), delta("nominal.")]), + doGenerate: (async () => ({})) as any, + params: params("status?") as any, + model: {} as any, + }); + expect(await drain(out.stream)).toBe("All systems nominal."); + }); + + it("blocks PII that completes across streamed deltas via wrapStream", async () => { + const mw = aegisMiddleware(); + const out = await mw.wrapStream!({ + doStream: async () => + streamOf([delta("Reach me at "), delta("admin@"), delta("example.com.")]), + doGenerate: (async () => ({})) as any, + params: params("contact?") as any, + model: {} as any, + }); + await expect(drain(out.stream)).rejects.toBeInstanceOf(AegisBlockedError); + }); });