feat(ai): extend the prompt injection defense to streamed suggestions

.eslintrc.js

@@ -10,7 +10,9 @@ module.exports = {
      * per .nvmrc) - not part of eslint's "node" env, which predates them
      */
     'ReadableStream': 'readonly',
-    'Response': 'readonly'
+    'Response': 'readonly',
+    'TransformStream': 'readonly',
+    'TransformStreamDefaultController': 'readonly'
   },
   rules: {
     '@typescript-eslint/camelcase': 'warn',

package.json

@@ -1,6 +1,6 @@
 {
   "name": "hawk.api",
-  "version": "1.5.12",
+  "version": "1.5.13",
   "main": "index.ts",
   "license": "BUSL-1.1",
   "scripts": {

src/integrations/vercel-ai/index.ts

@@ -1,5 +1,6 @@
-import { generateText, streamText } from 'ai';
+import { generateText, streamText, type TextStreamPart, type ToolSet } from 'ai';
 import { ProviderOptions } from '@ai-sdk/provider-utils';
+import type { GuardVerdict, StreamGuard } from '../../services/askAi/security/holdback';
 
 /**
  * Params for a single completion call to the model
@@ -16,6 +17,97 @@ export interface CompletionParams {
   prompt: string;
 }
 
+/**
+ * Params for a streaming completion call to the model
+ */
+export interface StreamParams extends CompletionParams {
+  /**
+   * Inspects the model's text before it leaves the server. Supplied by the
+   * service layer, because what counts as unsafe output is a domain question,
+   * not a transport one.
+   */
+  guard?: StreamGuard;
+
+  /**
+   * Called once, the first time the guard reports a leak
+   */
+  onLeak?: () => void;
+}
+
+/**
+ * Wrap the model's stream so every text delta passes through `guard`.
+ *
+ * Operates on typed stream parts rather than the encoded SSE bytes, where JSON
+ * envelopes and escaping would split a marker beyond the reach of any substring
+ * scan. The guard's holdback is released on `text-end`, so emitted deltas stay
+ * inside the text block they belong to; the TransformStream's own `flush` only
+ * covers a stream that ends without one.
+ *
+ * `stopStream` is not used: it obliges the caller to synthesize finish chunks
+ * whose shape follows the SDK version. Suppressing text keeps the stream well
+ * formed instead.
+ *
+ * @param guard - guard for this stream
+ * @param onLeak - called once when the guard first reports a leak
+ * @returns transform factory accepted by `streamText`
+ */
+function guardedTransform<TOOLS extends ToolSet>(guard: StreamGuard, onLeak?: () => void) {
+  return (): TransformStream<TextStreamPart<TOOLS>, TextStreamPart<TOOLS>> => {
+    let lastTextId: string | null = null;
+    let leakReported = false;
+
+    /**
+     * Forward the guard's verdict downstream, reporting a leak at most once
+     *
+     * @param verdict - what the guard allows to be sent
+     * @param controller - transform stream controller
+     * @param id - id of the text block the delta belongs to
+     */
+    const forward = (
+      verdict: GuardVerdict,
+      controller: TransformStreamDefaultController<TextStreamPart<TOOLS>>,
+      id: string | null
+    ): void => {
+      if (verdict.emit && id !== null) {
+        controller.enqueue({
+          type: 'text-delta',
+          id,
+          text: verdict.emit,
+        } as TextStreamPart<TOOLS>);
+      }
+
+      if (verdict.leaked && !leakReported) {
+        leakReported = true;
+
+        if (onLeak) {
+          onLeak();
+        }
+      }
+    };
+
+    return new TransformStream<TextStreamPart<TOOLS>, TextStreamPart<TOOLS>>({
+      transform(chunk, controller): void {
+        if (chunk.type === 'text-delta') {
+          lastTextId = chunk.id;
+          forward(guard.push(chunk.text), controller, chunk.id);
+
+          return;
+        }
+
+        if (chunk.type === 'text-end') {
+          forward(guard.flush(), controller, chunk.id);
+        }
+
+        controller.enqueue(chunk);
+      },
+
+      flush(controller): void {
+        forward(guard.flush(), controller, lastTextId);
+      },
+    });
+  };
+}
+
 /**
  * Interface for interacting with Vercel AI Gateway
  *
@@ -70,15 +162,17 @@ class VercelAIApi {
     /**
      * Send a system/prompt pair to the model and return the generated text as a stream
      *
-     * @param {CompletionParams} params - system instruction and prompt to complete
+     * @param {StreamParams} params - system instruction, prompt and optional output guard
      * @returns {StreamTextResult} text generated by the model, as a stream
      */
-    public stream({ system, prompt }: CompletionParams): ReturnType<typeof streamText> {
+    public stream({ system, prompt, guard, onLeak }: StreamParams): ReturnType<typeof streamText> {
       return streamText({
         model: this.modelId,
         system,
         prompt,
         providerOptions: this.providerOptions,
+        // eslint-disable-next-line camelcase, @typescript-eslint/camelcase
+        experimental_transform: guard ? guardedTransform(guard, onLeak) : undefined,
       });
     }
 }

src/services/ai.ts

@@ -2,6 +2,7 @@ import HawkCatcher from '@hawk.so/nodejs';
 import { vercelAIApi } from '../integrations/vercel-ai/';
 import { buildEventPrompt, spotlightInstruction } from './askAi/security/spotlighting';
 import { isLeaked, SUGGESTION_FALLBACK_MESSAGE } from './askAi/security/leakDetector';
+import { createLeakGuard } from './askAi/security/holdback';
 import { ctoInstruction } from './askAi/instructions/cto';
 import { EventsFactoryInterface } from './types';
 import type { Event } from './types';
@@ -70,8 +71,9 @@ export class AIService {
   /**
    * Generate streaming suggestion for the event
    *
-   * The payload is spotlighted by {@link buildEventPrompt} exactly as in
-   * {@link AIService.generateSuggestion}.
+   * Defended exactly as {@link AIService.generateSuggestion}, except that the
+   * answer is checked by {@link createLeakGuard} as it streams out rather than
+   * by {@link isLeaked} once it is complete.
    *
    * @param eventsFactory - events factory
    * @param eventId - event id
@@ -90,6 +92,8 @@ export class AIService {
     return vercelAIApi.stream({
       system: ctoInstruction + spotlightInstruction(nonce),
       prompt,
+      guard: createLeakGuard(nonce),
+      onLeak: () => reportRejectedSuggestion(eventId, originalEventId),
     });
   }
 

src/services/askAi/security/holdback.ts

@@ -0,0 +1,130 @@
+import { isLeaked, SUGGESTION_FALLBACK_MESSAGE } from './leakDetector';
+
+/**
+ * What the guard allows the transport to send downstream
+ */
+export interface GuardVerdict {
+  /**
+   * Text safe to forward now, which is what was fed in minus the holdback
+   */
+  emit: string;
+
+  /**
+   * Whether a leak was detected. Once true it stays true and no further model
+   * text is forwarded.
+   */
+  leaked: boolean;
+}
+
+/**
+ * Port implemented by the domain and consumed by the transport, so the
+ * provider adapter never needs to know what a leak is
+ */
+export interface StreamGuard {
+  /**
+   * Inspect the next piece of model output
+   *
+   * @param chunk - text delta produced by the model
+   * @returns {GuardVerdict} text safe to forward now
+   */
+  push(chunk: string): GuardVerdict;
+
+  /**
+   * Release whatever is still withheld, at end of stream
+   *
+   * @returns {GuardVerdict} remaining text safe to forward
+   */
+  flush(): GuardVerdict;
+}
+
+/**
+ * Streaming counterpart of {@link isLeaked}.
+ *
+ * The nonce can arrive split across two deltas, so scanning each delta alone
+ * would never see it whole. The guard therefore keeps a *holdback*: the last
+ * `nonce.length - 1` characters fed in so far, kept unsent. Every new delta is
+ * scanned together with the holdback, and only the part that can no longer
+ * begin the nonce is released.
+ *
+ * That length is the exact minimum. An occurrence of the nonce spans
+ * `nonce.length` characters, so holding one less guarantees it falls inside a
+ * single scanned window and never reaches the client.
+ *
+ * On detection the rest of the answer is replaced by
+ * {@link SUGGESTION_FALLBACK_MESSAGE}, emitted once. The prefix already sent
+ * cannot be retracted, which is acceptable: the nonce is what triggered
+ * detection and is still held back when it fires.
+ *
+ * @param nonce - per-request nonce used in the prompt markers
+ * @returns {StreamGuard} guard for a single stream, not reusable
+ */
+export function createLeakGuard(nonce: string): StreamGuard {
+  const holdback = Math.max(nonce.length - 1, 0);
+
+  let withheld = '';
+  let leaked = false;
+
+  /**
+   * Mark the stream as leaked and produce the one verdict that still carries
+   * text: the fallback message
+   *
+   * @returns {GuardVerdict} verdict replacing the rest of the answer
+   */
+  const reject = (): GuardVerdict => {
+    leaked = true;
+    withheld = '';
+
+    return {
+      emit: SUGGESTION_FALLBACK_MESSAGE,
+      leaked: true,
+    };
+  };
+
+  return {
+    push(chunk: string): GuardVerdict {
+      if (leaked) {
+        return {
+          emit: '',
+          leaked: true,
+        };
+      }
+
+      const window = withheld + chunk;
+
+      if (isLeaked(window, nonce)) {
+        return reject();
+      }
+
+      const sendable = Math.max(window.length - holdback, 0);
+
+      withheld = window.slice(sendable);
+
+      return {
+        emit: window.slice(0, sendable),
+        leaked: false,
+      };
+    },
+
+    flush(): GuardVerdict {
+      if (leaked) {
+        return {
+          emit: '',
+          leaked: true,
+        };
+      }
+
+      const pending = withheld;
+
+      withheld = '';
+
+      if (isLeaked(pending, nonce)) {
+        return reject();
+      }
+
+      return {
+        emit: pending,
+        leaked: false,
+      };
+    },
+  };
+}