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

.eslintrc.js

@@ -4,6 +4,13 @@ module.exports = {
     'node': true,
     'jest': true
   },
+  globals: {
+    /**
+     * Global since Node 18 (this project runs Node 24 per .nvmrc), but not part of
+     * eslint's "node" env, which predates the WHATWG Streams API
+     */
+    'TransformStream': 'readonly'
+  },
   rules: {
     '@typescript-eslint/camelcase': 'warn',
     '@typescript-eslint/no-unused-vars': 'warn',

package.json

@@ -1,6 +1,6 @@
 {
   "name": "hawk.api",
-  "version": "1.5.14",
+  "version": "1.5.15",
   "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,108 @@ 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. Required: a stream cannot be checked after the fact,
+   * so an omitted guard would mean an unchecked answer.
+   */
+  guard: StreamGuard;
+
+  /**
+   * Called once, the first time the guard rejects the answer
+   */
+  onReject: () => 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 onReject - called once when the guard first rejects the answer
+ * @returns transform factory accepted by `streamText`
+ */
+function guardedTransform<TOOLS extends ToolSet>(guard: StreamGuard, onReject: () => void) {
+  return (): TransformStream<TextStreamPart<TOOLS>, TextStreamPart<TOOLS>> => {
+    let lastTextId: string | null = null;
+    let rejectReported = false;
+
+    /**
+     * Forward the guard's verdict downstream, reporting a rejection at most once.
+     *
+     * A rejection travels as an error part rather than more text, so the client
+     * can tell it apart from the answer and drop what it has already rendered.
+     *
+     * @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.rejected) {
+        if (!rejectReported) {
+          rejectReported = true;
+
+          controller.enqueue({
+            type: 'error',
+            error: new Error(verdict.emit),
+          } as TextStreamPart<TOOLS>);
+
+          onReject();
+        }
+
+        return;
+      }
+
+      if (verdict.emit && id !== null) {
+        controller.enqueue({
+          type: 'text-delta',
+          id,
+          text: verdict.emit,
+        } as TextStreamPart<TOOLS>);
+      }
+    };
+
+    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
  *
@@ -68,15 +171,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 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, onReject }: StreamParams): ReturnType<typeof streamText> {
     return streamText({
       model: this.modelId,
       system,
       prompt,
       providerOptions: this.providerOptions,
+      // eslint-disable-next-line camelcase, @typescript-eslint/camelcase
+      experimental_transform: guardedTransform(guard, onReject),
     });
   }
 }

src/services/askAi/routes.ts

@@ -95,7 +95,8 @@ export function createAiStreamRouter(): express.Router {
         return;
       }
 
-      result.pipeUIMessageStreamToResponse(res);
+      /** Reasoning is neither scanned for the nonce nor rendered anywhere */
+      result.pipeUIMessageStreamToResponse(res, { sendReasoning: false });
     } catch (error) {
       next(error);
     }

src/services/askAi/security/holdback.ts

@@ -0,0 +1,146 @@
+import { echoesNonce, SUGGESTION_FALLBACK_MESSAGE } from './nonceEcho';
+
+/**
+ * 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 the nonce was found. Once true it stays true and no further model
+   * text is forwarded.
+   */
+  rejected: boolean;
+}
+
+/**
+ * Port implemented by the domain and consumed by the transport, so the
+ * provider adapter never needs to know what makes output unsafe
+ */
+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 the end of a text block
+   *
+   * @returns {GuardVerdict} remaining text safe to forward
+   */
+  flush(): GuardVerdict;
+}
+
+/**
+ * Streaming counterpart of {@link echoesNonce}.
+ *
+ * The nonce can arrive split across two deltas, so scanning each delta alone
+ * would never see it whole. The guard 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 spans `nonce.length`
+ * characters, so holding one less leaves it inside a single scanned window.
+ *
+ * {@link StreamGuard.flush} has to release the holdback at the end of a text
+ * block, because that text belongs to the block and cannot be emitted under the
+ * next one's id. Its tail is kept as scanning context instead, without which a
+ * nonce split across two blocks would pass unseen.
+ *
+ * On rejection nothing more is released and {@link SUGGESTION_FALLBACK_MESSAGE}
+ * is returned once, for the transport to deliver as it sees fit. Text already
+ * sent cannot be taken back, and the holdback cuts it at an arbitrary character.
+ *
+ * @param nonce - per-request nonce used in the prompt markers
+ * @returns {StreamGuard} guard for a single stream, not reusable
+ */
+export function createStreamGuard(nonce: string): StreamGuard {
+  const holdback = Math.max(nonce.length - 1, 0);
+
+  let withheld = '';
+  let sentTail = '';
+  let rejected = false;
+
+  /**
+   * Keep only as much already-sent text as a nonce could still overlap
+   *
+   * @param text - text sent so far, ending with what was just emitted
+   * @returns {string} trailing scanning context
+   */
+  const keepTail = (text: string): string => text.slice(Math.max(text.length - holdback, 0));
+
+  /**
+   * Mark the stream as rejected and produce the one verdict that still carries
+   * text: the fallback message
+   *
+   * @returns {GuardVerdict} verdict replacing the rest of the answer
+   */
+  const reject = (): GuardVerdict => {
+    rejected = true;
+    withheld = '';
+    sentTail = '';
+
+    return {
+      emit: SUGGESTION_FALLBACK_MESSAGE,
+      rejected: true,
+    };
+  };
+
+  return {
+    push(chunk: string): GuardVerdict {
+      if (rejected) {
+        return {
+          emit: '',
+          rejected: true,
+        };
+      }
+
+      if (echoesNonce(sentTail + withheld + chunk, nonce)) {
+        return reject();
+      }
+
+      const pending = withheld + chunk;
+      const sendable = Math.max(pending.length - holdback, 0);
+      const emit = pending.slice(0, sendable);
+
+      withheld = pending.slice(sendable);
+      sentTail = keepTail(sentTail + emit);
+
+      return {
+        emit,
+        rejected: false,
+      };
+    },
+
+    flush(): GuardVerdict {
+      if (rejected) {
+        return {
+          emit: '',
+          rejected: true,
+        };
+      }
+
+      const pending = withheld;
+
+      withheld = '';
+
+      if (echoesNonce(sentTail + pending, nonce)) {
+        return reject();
+      }
+
+      sentTail = keepTail(sentTail + pending);
+
+      return {
+        emit: pending,
+        rejected: false,
+      };
+    },
+  };
+}

src/services/askAi/service.ts

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