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