feat(ai): guard streamed output in the Vercel AI SDK middleware (wrapStream)

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 };
+    },
   };
 }

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<any>): Promise<string> {
+    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);
+  });
 });