From e64c1a7896a969c502d36cedddbd3fb8ef73d5d5 Mon Sep 17 00:00:00 2001 From: Reversean Date: Wed, 29 Jul 2026 19:36:35 +0300 Subject: [PATCH] feat(ai): extend the prompt injection defense to streamed suggestions The one-shot path scans the finished answer before returning it. A streamed answer leaves the server while it is still being written, so the same scan cannot be applied: a nonce split across two deltas passes a per-delta check untouched and the marker reaches the client. The stream now runs through a guard holding back the last nonce.length - 1 characters and scanning them together with each new delta, releasing only text that can no longer begin the nonce. That length is the exact minimum for an occurrence to always fall inside a single scanned window. The guard is applied to the model's typed stream parts rather than the encoded SSE bytes, where JSON envelopes and escaping would split the nonce beyond the reach of a substring scan. On detection the rest of the answer is replaced by the fallback message and the event ids are logged. The stream is not aborted: aborting obliges the caller to synthesize finish chunks whose shape follows the SDK version, and suppressing text keeps the stream well formed instead. What counts as a leak stays in the domain layer behind a port, so the provider adapter still knows nothing about it. --- .eslintrc.js | 4 +- package.json | 2 +- src/integrations/vercel-ai/index.ts | 100 +++++++++++- src/services/ai.ts | 8 +- src/services/askAi/security/holdback.ts | 130 +++++++++++++++ test/integrations/vercel-ai.test.ts | 209 ++++++++++++++++++++++++ test/services/askAi-holdback.test.ts | 87 ++++++++++ test/services/askAi.test.ts | 26 +++ 8 files changed, 559 insertions(+), 7 deletions(-) create mode 100644 src/services/askAi/security/holdback.ts create mode 100644 test/services/askAi-holdback.test.ts diff --git a/.eslintrc.js b/.eslintrc.js index 7cec6e3f..30294544 100644 --- a/.eslintrc.js +++ b/.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', diff --git a/package.json b/package.json index 638a7e30..db9cff2d 100644 --- a/package.json +++ b/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": { diff --git a/src/integrations/vercel-ai/index.ts b/src/integrations/vercel-ai/index.ts index 88812e7e..16182f55 100644 --- a/src/integrations/vercel-ai/index.ts +++ b/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(guard: StreamGuard, onLeak?: () => void) { + return (): TransformStream, TextStreamPart> => { + 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>, + id: string | null + ): void => { + if (verdict.emit && id !== null) { + controller.enqueue({ + type: 'text-delta', + id, + text: verdict.emit, + } as TextStreamPart); + } + + if (verdict.leaked && !leakReported) { + leakReported = true; + + if (onLeak) { + onLeak(); + } + } + }; + + return new TransformStream, TextStreamPart>({ + 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 { + public stream({ system, prompt, guard, onLeak }: StreamParams): ReturnType { 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, }); } } diff --git a/src/services/ai.ts b/src/services/ai.ts index ed9d689b..0d8c1b6d 100644 --- a/src/services/ai.ts +++ b/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), }); } diff --git a/src/services/askAi/security/holdback.ts b/src/services/askAi/security/holdback.ts new file mode 100644 index 00000000..bc8537e7 --- /dev/null +++ b/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, + }; + }, + }; +} diff --git a/test/integrations/vercel-ai.test.ts b/test/integrations/vercel-ai.test.ts index 624b26d8..e73f3022 100644 --- a/test/integrations/vercel-ai.test.ts +++ b/test/integrations/vercel-ai.test.ts @@ -1,12 +1,110 @@ import '../../src/env-test'; import { generateText, streamText } from 'ai'; import { vercelAIApi } from '../../src/integrations/vercel-ai/'; +import type { GuardVerdict, StreamGuard } from '../../src/services/askAi/security/holdback'; jest.mock('ai', () => ({ generateText: jest.fn(), streamText: jest.fn(), })); +const testTextId = 'text-block-1'; + +/** + * Build a text-delta chunk as the SDK would produce it + * + * @param text - delta text + * @returns text-delta stream part + */ +function delta(text: string): Record { + return { + type: 'text-delta', + id: testTextId, + text, + }; +} + +/** + * Pull the transform factory the transport handed to streamText + * + * @param mock - mocked streamText + * @returns the experimental_transform argument of the last call + */ +function transformOf(mock: jest.Mock): unknown { + return mock.mock.calls[mock.mock.calls.length - 1][0].experimental_transform; +} + +/** + * Drive the transport's real transform: ask it for a stream by calling + * `stream` with a guard, then push chunks through what it registered. + * + * @param guard - guard to install + * @param chunks - stream parts to feed + * @param onLeak - leak callback to install + * @returns every stream part the transform let out + */ +async function runTransformRaw( + guard: StreamGuard, + chunks: Record[], + onLeak?: () => void +): Promise<{ parts: Record[] }> { + (streamText as jest.Mock).mockReturnValue({}); + + vercelAIApi.stream({ + system: 'system', + prompt: 'prompt', + guard, + onLeak, + }); + + const factory = transformOf(streamText as jest.Mock) as () => TransformStream; + const transform = factory(); + const writer = transform.writable.getWriter(); + const reader = transform.readable.getReader(); + const parts: Record[] = []; + + const collecting = (async (): Promise => { + for (;;) { + const { done, value } = await reader.read(); + + if (done) { + return; + } + + parts.push(value as Record); + } + })(); + + for (const chunk of chunks) { + await writer.write(chunk); + } + + await writer.close(); + await collecting; + + return { parts }; +} + +/** + * Same as `runTransformRaw`, reduced to the text that reached the client + * + * @param guard - guard to install + * @param chunks - stream parts to feed + * @param onLeak - leak callback to install + * @returns text of every text-delta the transform let out + */ +async function runTransform( + guard: StreamGuard, + chunks: Record[], + onLeak?: () => void +): Promise { + const { parts } = await runTransformRaw(guard, chunks, onLeak); + + return parts + .filter((part) => part.type === 'text-delta') + .map((part) => part.text as string); +} + describe('VercelAIApi', () => { const testSystem = 'system instruction'; const testPrompt = 'user prompt'; @@ -59,5 +157,116 @@ describe('VercelAIApi', () => { }); expect(result).toBe(streamResult); }); + + it('should not install a stream transform when no guard is supplied', () => { + (streamText as jest.Mock).mockReturnValue({}); + + vercelAIApi.stream({ + system: testSystem, + prompt: testPrompt, + }); + + expect(transformOf(streamText as jest.Mock)).toBeUndefined(); + }); + }); + + describe('stream guard wiring', () => { + /** + * Stand-in for the real leak guard: records what it was fed and dictates + * what may leave, so the transport's plumbing is tested without the + * detection logic (covered by askAi-holdback.test.ts) + * + * @param verdicts - verdicts to return, in order; anything past the end + * passes the text through unchanged + * @returns {StreamGuard} guard stub that also exposes what it was fed + */ + const stubGuard = (verdicts: GuardVerdict[]): StreamGuard & { fed: string[] } => { + const fed: string[] = []; + let call = 0; + + return { + fed, + push: (chunk: string): GuardVerdict => { + fed.push(chunk); + + return verdicts[call++] ?? { + emit: chunk, + leaked: false, + }; + }, + flush: (): GuardVerdict => verdicts[call++] ?? { + emit: '', + leaked: false, + }, + }; + }; + + it('should route text deltas through the guard and forward only what it allows', async () => { + const guard = stubGuard([ + { + emit: 'allowed', + leaked: false, + }, + { + emit: '', + leaked: false, + }, + ]); + + const emitted = await runTransform(guard, [delta('first'), delta('second')]); + + expect(guard.fed).toEqual(['first', 'second']); + expect(emitted).toEqual([ 'allowed' ]); + }); + + it('should release the withheld tail when the model closes the text block', async () => { + const guard = stubGuard([ + { + emit: '', + leaked: false, + }, + { + emit: 'tail', + leaked: false, + }, + ]); + + const emitted = await runTransform(guard, [delta('held back'), { + type: 'text-end', + id: testTextId, + } ]); + + expect(emitted).toEqual([ 'tail' ]); + }); + + it('should report a leak once even when the guard keeps reporting it', async () => { + const onLeak = jest.fn(); + const guard = stubGuard([ + { + emit: '', + leaked: true, + }, + { + emit: '', + leaked: true, + }, + ]); + + await runTransform(guard, [delta('first'), delta('second')], onLeak); + + expect(onLeak).toHaveBeenCalledTimes(1); + }); + + it('should pass non-text chunks through untouched', async () => { + const guard = stubGuard([]); + const chunks = [ { + type: 'text-start', + id: testTextId, + } ]; + + const { parts } = await runTransformRaw(guard, chunks); + + expect(parts).toContainEqual(chunks[0]); + }); }); }); diff --git a/test/services/askAi-holdback.test.ts b/test/services/askAi-holdback.test.ts new file mode 100644 index 00000000..26639e7d --- /dev/null +++ b/test/services/askAi-holdback.test.ts @@ -0,0 +1,87 @@ +import '../../src/env-test'; +import { createLeakGuard } from '../../src/services/askAi/security/holdback'; +import { SUGGESTION_FALLBACK_MESSAGE } from '../../src/services/askAi/security/leakDetector'; + +const nonce = '0123456789abcdef0123456789abcdef'; + +/** + * Feed a whole answer through a guard one chunk at a time + * + * @param chunks - text deltas as the model would produce them + * @returns {object} text the client would have received, and whether a leak fired + */ +function drain(chunks: string[]): { emitted: string; leaked: boolean } { + const guard = createLeakGuard(nonce); + let emitted = ''; + let leaked = false; + + for (const chunk of chunks) { + const verdict = guard.push(chunk); + + emitted += verdict.emit; + leaked = leaked || verdict.leaked; + } + + const final = guard.flush(); + + return { + emitted: emitted + final.emit, + leaked: leaked || final.leaked, + }; +} + +describe('createLeakGuard', () => { + it('should pass a clean answer through unchanged', () => { + const chunks = ['## Cause\n', 'The variable is not defined. ', 'Check the initialisation.']; + + expect(drain(chunks)).toEqual({ + emitted: chunks.join(''), + leaked: false, + }); + }); + + it('should withhold the tail until enough text has arrived to clear it', () => { + const guard = createLeakGuard(nonce); + const answer = 'short'; + + expect(guard.push(answer).emit).toBe(''); + expect(guard.flush().emit).toBe(answer); + }); + + it('should detect a nonce split across two chunks without ever emitting it', () => { + const result = drain([`marker ${nonce.slice(0, 20)}`, `${nonce.slice(20)} tail`]); + + expect(result.leaked).toBe(true); + expect(result.emitted).not.toContain(nonce); + }); + + it('should detect the nonce echoed in a different case', () => { + expect(drain([ `marker ${nonce.toUpperCase()} tail` ]).leaked).toBe(true); + }); + + it('should stay silent for the rest of the stream after a leak', () => { + const guard = createLeakGuard(nonce); + + guard.push(`marker ${nonce}`); + + expect(guard.push('the rest of the answer')).toEqual({ + emit: '', + leaked: true, + }); + expect(guard.flush()).toEqual({ + emit: '', + leaked: true, + }); + }); + + it('should replace the rest of the answer with the fallback message exactly once', () => { + const guard = createLeakGuard(nonce); + const emissions = [ + guard.push(`marker ${nonce}`).emit, + guard.push('more text').emit, + guard.flush().emit, + ].filter(Boolean); + + expect(emissions).toEqual([ SUGGESTION_FALLBACK_MESSAGE ]); + }); +}); diff --git a/test/services/askAi.test.ts b/test/services/askAi.test.ts index 2c80cd8f..9618f733 100644 --- a/test/services/askAi.test.ts +++ b/test/services/askAi.test.ts @@ -3,6 +3,7 @@ import HawkCatcher from '@hawk.so/nodejs'; import { EventAddons, EventData } from '@hawk.so/types'; import { AIService } from '../../src/services/ai'; import { vercelAIApi } from '../../src/integrations/vercel-ai/'; +import type { StreamParams } from '../../src/integrations/vercel-ai/'; import { ctoInstruction } from '../../src/services/askAi/instructions/cto'; import { UNTRUSTED_DATA_MARKER_NAME } from '../../src/services/askAi/security/spotlighting'; import { SUGGESTION_FALLBACK_MESSAGE } from '../../src/services/askAi/security/leakDetector'; @@ -151,5 +152,30 @@ describe('AIService', () => { expect(vercelAIApi.stream).not.toHaveBeenCalled(); }); + + it('should arm the guard with this request\'s nonce and report the event ids when it fires', async () => { + (vercelAIApi.stream as jest.Mock).mockReturnValue({}); + + await aiService.streamSuggestion(eventsFactoryWithPayload(), testEventId, testOriginalEventId); + + const args = (vercelAIApi.stream as jest.Mock).mock.calls[0][0] as StreamParams; + const eventIds = expect.objectContaining({ + eventId: testEventId, + originalEventId: testOriginalEventId, + }); + + /** + * A guard built from any other nonce would let this stream's marker + * through, so feeding it this request's nonce must trip it + */ + expect(args.guard?.push(`service marker ${nonceFromPrompt(args.prompt)}`).leaked).toBe(true); + + if (args.onLeak) { + args.onLeak(); + } + + expect(consoleErrorSpy).toHaveBeenCalledWith(expect.any(String), eventIds); + expect(HawkCatcher.send).toHaveBeenCalledWith(expect.any(Error), eventIds); + }); }); });