@intflows/genkit-guard 0.0.9 → 0.0.11

package/dist/middleware/middleware.js

@@ -2,6 +2,8 @@ import { generateMiddleware, z } from 'genkit';
 import { analyzeIntentStructured, detectInjection } from '../intent/intentAnalyzer.js';
 import { detectPII } from '../pii/detector.js';
 import { PiiTokenizer } from '../pii/tokenizer.js';
+import { defaultPiiVaultStorage } from '../pii/storage.js';
+const GUARD_CONTEXT_KEY = '__genkitGuard';
 const guardConfigSchema = z.object({
     intent: z.object({
         mode: z.string().optional(),
@@ -15,6 +17,15 @@ const guardConfigSchema = z.object({
         reversible: z.boolean().optional(),
         model: z.string().optional(),
         mode: z.enum(['ner', 'classifier']).optional(),
+        vault: z.object({
+            storage: z.any().optional(),
+            scopeId: z.any().optional(),
+        }).optional(),
+    }).optional(),
+    logging: z.object({
+        enabled: z.boolean().optional(),
+        level: z.enum(['debug', 'info', 'warn', 'error']).optional(),
+        serviceName: z.string().optional(),
     }).optional(),
     models: z.object({
         extractor: z.string().optional(),
@@ -22,83 +33,132 @@ const guardConfigSchema = z.object({
 }).passthrough();
 export const guardMiddleware = generateMiddleware({
     name: 'genkitGuard',
-    description: 'Blocks prompt injection and disallowed intent, then masks PII before model calls and unmasks model responses.',
+    description: 'Blocks prompt injection and disallowed intent, masks PII before model calls, restores PII for tool calls, and audits tool PII access.',
     configSchema: guardConfigSchema,
 }, ({ config }) => createGuardHooks(config));
 export const guardPlugin = guardMiddleware.plugin;
 export function guard(config) {
     const hooks = createGuardHooks(config);
     const baseMiddleware = guardMiddleware(config);
-    // 1. Create the wrapper function runner
     const fnRunner = async (req, ctxOrNext, maybeNext) => {
         if (typeof maybeNext === 'function') {
             return hooks.model(req, ctxOrNext, maybeNext);
         }
         return hooks.model(req, {}, async (modifiedReq) => ctxOrNext(modifiedReq || req));
     };
-    // 2. Combine the base middleware properties and custom hooks into a source object
     const source = Object.assign({}, baseMiddleware, hooks);
-    // 3. Safely copy properties onto the function runner, explicitly skipping the read-only 'name' property
     for (const key of Object.keys(source)) {
         if (key === 'name')
-            continue; // Prevent the TypeError
-        // Use defineProperty or simple assignment for everything else
+            continue;
         Object.defineProperty(fnRunner, key, {
             value: source[key],
             writable: true,
             configurable: true,
-            enumerable: true
+            enumerable: true,
         });
     }
     return fnRunner;
 }
+export const guardAction = guard;
 function createGuardHooks(config) {
+    const logger = createLogger(config);
     return {
         model: async (req, ctx, next) => {
             const input = getInputText(req);
+            logger('info', 'guard.model.start', 'Starting guard checks for model request');
             const isInjection = await detectInjection(input);
             if (isInjection) {
-                console.warn('[Intent Guard] Prompt injection pattern detected in input');
+                logger('warn', 'guard.intent.blocked', 'Prompt injection pattern detected', {
+                    reason: 'pattern_match',
+                });
                 return block('Prompt injection detected', {
                     reason: 'pattern_match',
                 });
             }
-            console.log("[Intent Guard] Analyzing intent for user's input");
+            logger('info', 'guard.intent.analysis.start', 'Analyzing request intent');
             const intentResult = await analyzeIntentStructured(input, config?.intent?.semantic?.intents ?? {}, config?.intent?.semantic?.threshold ?? 0.7);
-            console.log(`[Intent Guard] Detected intent: ${intentResult.intent} (score: ${intentResult.score.toFixed(2)})`);
+            logger('info', 'guard.intent.analysis.complete', 'Intent analysis completed', {
+                intent: intentResult.intent,
+                score: roundScore(intentResult.score),
+                allowed: intentResult.allowed,
+            });
             if (!intentResult.allowed) {
-                console.warn(`[Intent Guard] Intent "${intentResult.intent}" not allowed ${intentResult.allowed}`);
+                logger('warn', 'guard.intent.blocked', 'Intent not allowed', {
+                    intent: intentResult.intent,
+                    score: roundScore(intentResult.score),
+                });
                 return block('Intent not allowed', {
                     intent: intentResult.intent,
                     score: intentResult.score,
                 });
             }
-            const piiResponse = await detectPII(input, {
-                model: config?.pii?.model,
-                mode: config?.pii?.mode,
-            });
+            const textForPii = collectModelRequestText(req);
+            const piiResponse = await scanPII(textForPii, config);
             const piiMatches = piiResponse?.matches || [];
-            console.log(`[PII Guard] Detected PII: ${piiMatches.length} matches found` + (piiResponse.classifier ? ' (classifier output present)' : ''));
-            const tokenizer = new PiiTokenizer();
-            const piiResult = tokenizer.mask(input, piiMatches);
-            console.log(`[PII Guard] Masked PII: ${piiResult.piiTypes.length} types found`);
+            const tokenizer = createTokenizer(config, req, ctx);
+            const piiTypes = uniqueTypes(piiMatches);
+            await tokenizer.importTokens(textForPii);
+            await maskModelRequest(req, tokenizer, piiMatches);
+            pushTokenizer(ctx, tokenizer);
+            logger(piiMatches.length > 0 ? 'warn' : 'info', 'guard.model.pii.masked', 'PII scan completed for model request', {
+                piiDetected: piiMatches.length > 0,
+                piiMatchCount: piiMatches.length,
+                piiTypes,
+                piiMode: config?.pii?.mode ?? 'ner',
+                classifierOutputPresent: Boolean(piiResponse.classifier),
+            });
             req.metadata = {
                 ...req.metadata,
-                piiTokenizer: tokenizer,
                 intent: intentResult.intent,
                 score: intentResult.score,
                 piiDetected: piiMatches.length > 0,
-                piiTypes: piiResult.piiTypes,
-                maskedInput: piiResult.maskedText,
+                piiTypes,
+                maskedInput: getInputText(req),
                 piiModel: config?.pii?.model,
                 piiMode: config?.pii?.mode,
                 piiClassifierOutput: piiResponse.classifier,
             };
-            replaceInputText(req, piiResult.maskedText);
             const res = await next(req, ctx);
-            console.log('[PII Guard] Unmasking response if needed');
-            unmaskResponse(res, tokenizer);
-            console.log('[PII Guard] Deep unmasking complete across all candidates and custom fields.');
+            const unmaskedResponse = await unmaskObject(res, getGuardState(ctx).tokenizers);
+            logger('info', 'guard.model.response.unmasked', 'Model response unmasked for downstream execution', {
+                piiTypes,
+            });
+            return unmaskedResponse;
+        },
+        tool: async (req, ctx, next) => {
+            const state = getGuardState(ctx);
+            const toolName = req?.toolRequest?.name;
+            if (req?.toolRequest && 'input' in req.toolRequest) {
+                req.toolRequest.input = await unmaskObject(req.toolRequest.input, state.tokenizers);
+            }
+            const toolInputText = collectStrings(req?.toolRequest?.input).join('\n');
+            const piiResponse = await scanPII(toolInputText, config);
+            const piiMatches = piiResponse?.matches || [];
+            const piiTypes = uniqueTypes(piiMatches);
+            req.metadata = {
+                ...req.metadata,
+                piiDetected: piiMatches.length > 0,
+                piiTypes,
+                piiMatchCount: piiMatches.length,
+            };
+            logger(piiMatches.length > 0 ? 'warn' : 'info', 'guard.tool.pii.checked', 'Tool request PII scan completed', {
+                toolName,
+                piiDetected: piiMatches.length > 0,
+                piiMatchCount: piiMatches.length,
+                piiTypes,
+            });
+            const res = await next(req, ctx);
+            const toolResponseText = collectStrings(res).join('\n');
+            if (toolResponseText) {
+                const responsePii = await scanPII(toolResponseText, config);
+                const responseMatches = responsePii?.matches || [];
+                logger(responseMatches.length > 0 ? 'warn' : 'info', 'guard.tool.response.pii.checked', 'Tool response PII scan completed', {
+                    toolName,
+                    piiDetected: responseMatches.length > 0,
+                    piiMatchCount: responseMatches.length,
+                    piiTypes: uniqueTypes(responseMatches),
+                });
+            }
             return res;
         },
     };
@@ -115,36 +175,149 @@ function getInputText(req) {
     if (typeof firstContent === 'string') {
         return firstContent;
     }
-    return '';
+    return collectStrings(lastMessage).join('\n');
 }
-function replaceInputText(req, text) {
+function collectModelRequestText(req) {
+    return [
+        ...collectStrings(req?.prompt),
+        ...collectStrings(req?.messages),
+        ...collectStrings(req?.docs),
+    ].join('\n');
+}
+async function maskModelRequest(req, tokenizer, matches) {
     if (typeof req.prompt === 'string') {
-        req.prompt = text;
+        req.prompt = (await tokenizer.mask(req.prompt, matches)).maskedText;
+    }
+    if (req.messages) {
+        req.messages = await transformStrings(req.messages, async (value) => (await tokenizer.mask(value, matches)).maskedText);
+    }
+    if (req.docs) {
+        req.docs = await transformStrings(req.docs, async (value) => (await tokenizer.mask(value, matches)).maskedText);
     }
-    req.messages = [
-        {
-            role: 'user',
-            content: [{ text }],
-        },
-    ];
 }
-function unmaskResponse(res, tokenizer) {
-    const transform = (obj) => {
-        if (typeof obj === 'string') {
-            return tokenizer.unmask(obj);
+async function unmaskObject(obj, tokenizers) {
+    return transformStrings(obj, async (value) => {
+        let result = value;
+        for (const tokenizer of tokenizers) {
+            result = await tokenizer.unmask(result);
         }
-        if (Array.isArray(obj)) {
-            return obj.map(transform);
+        return result;
+    });
+}
+async function transformStrings(obj, transform) {
+    if (typeof obj === 'string') {
+        return await transform(obj);
+    }
+    if (Array.isArray(obj)) {
+        for (let i = 0; i < obj.length; i++) {
+            obj[i] = await transformStrings(obj[i], transform);
         }
-        if (obj !== null && typeof obj === 'object') {
-            for (const key of Object.keys(obj)) {
-                obj[key] = transform(obj[key]);
-            }
-            return obj;
+        return obj;
+    }
+    if (obj !== null && typeof obj === 'object') {
+        for (const key of Object.keys(obj)) {
+            obj[key] = await transformStrings(obj[key], transform);
         }
         return obj;
+    }
+    return obj;
+}
+function collectStrings(obj) {
+    if (typeof obj === 'string') {
+        return [obj];
+    }
+    if (Array.isArray(obj)) {
+        return obj.flatMap(collectStrings);
+    }
+    if (obj !== null && typeof obj === 'object') {
+        return Object.values(obj).flatMap(collectStrings);
+    }
+    return [];
+}
+async function scanPII(text, config) {
+    if (!text.trim()) {
+        return {
+            matches: [],
+            classifier: undefined,
+        };
+    }
+    return detectPII(text, {
+        model: config?.pii?.model,
+        mode: config?.pii?.mode,
+    });
+}
+function getGuardState(ctx = {}) {
+    ctx.context = ctx.context || {};
+    ctx.context[GUARD_CONTEXT_KEY] = ctx.context[GUARD_CONTEXT_KEY] || { tokenizers: [] };
+    return ctx.context[GUARD_CONTEXT_KEY];
+}
+function pushTokenizer(ctx, tokenizer) {
+    const state = getGuardState(ctx);
+    state.tokenizers.push(tokenizer);
+}
+function createTokenizer(config, req, ctx) {
+    const configuredScope = config?.pii?.vault?.scopeId;
+    const scopeId = typeof configuredScope === 'function'
+        ? configuredScope(req, ctx)
+        : configuredScope;
+    return new PiiTokenizer({
+        scopeId: typeof scopeId === 'string' ? scopeId : undefined,
+        storage: config?.pii?.vault?.storage ?? defaultPiiVaultStorage,
+    });
+}
+function uniqueTypes(matches) {
+    return Array.from(new Set(matches.map((match) => match.type.toLowerCase())));
+}
+function roundScore(score) {
+    return Math.round(score * 10000) / 10000;
+}
+function createLogger(config) {
+    const enabled = config?.logging?.enabled ?? true;
+    const minimumLevel = config?.logging?.level ?? 'info';
+    const serviceName = config?.logging?.serviceName ?? '@intflows/genkit-guard';
+    const levelRank = {
+        debug: 10,
+        info: 20,
+        warn: 30,
+        error: 40,
+    };
+    const severityNumber = {
+        debug: 5,
+        info: 9,
+        warn: 13,
+        error: 17,
+    };
+    return (severity, eventName, body, attributes = {}) => {
+        if (!enabled || levelRank[severity] < levelRank[minimumLevel]) {
+            return;
+        }
+        const record = {
+            timestamp: new Date().toISOString(),
+            severityText: severity.toUpperCase(),
+            severityNumber: severityNumber[severity],
+            body,
+            resource: {
+                attributes: {
+                    'service.name': serviceName,
+                },
+            },
+            attributes: {
+                'event.name': eventName,
+                'code.namespace': 'genkit-guard',
+                ...attributes,
+            },
+        };
+        const line = JSON.stringify(record);
+        if (severity === 'error') {
+            console.error(line);
+        }
+        else if (severity === 'warn') {
+            console.warn(line);
+        }
+        else {
+            console.log(line);
+        }
     };
-    transform(res);
 }
 function block(message, metadata) {
     return {

package/dist/pii/storage.d.ts

@@ -0,0 +1,42 @@
+export type PiiVaultEntry = {
+    token: string;
+    value: string;
+};
+export interface PiiVaultStorage {
+    get(scopeId: string, token: string): string | undefined | Promise<string | undefined>;
+    getByToken?(token: string): string | undefined | Promise<string | undefined>;
+    set(scopeId: string, token: string, value: string): void | Promise<void>;
+    entries(scopeId: string): PiiVaultEntry[] | Promise<PiiVaultEntry[]>;
+}
+export type PiiVaultStorageAdapter = {
+    get: PiiVaultStorage['get'];
+    getByToken?: PiiVaultStorage['getByToken'];
+    set: PiiVaultStorage['set'];
+    entries: PiiVaultStorage['entries'];
+};
+export declare function createPiiVaultStorage(adapter: PiiVaultStorageAdapter): PiiVaultStorage;
+export type RedisPiiVaultClient = {
+    hGet?: (key: string, field: string) => Promise<string | null | undefined>;
+    hSet?: (key: string, field: string, value: string) => Promise<unknown>;
+    hGetAll?: (key: string) => Promise<Record<string, string>>;
+    hget?: (key: string, field: string) => Promise<string | null | undefined>;
+    hset?: (key: string, field: string, value: string) => Promise<unknown>;
+    hgetall?: (key: string) => Promise<Record<string, string>>;
+    expire?: (key: string, seconds: number) => Promise<unknown>;
+};
+export type RedisPiiVaultStorageOptions = {
+    keyPrefix?: string;
+    tokenIndexKey?: string;
+    ttlSeconds?: number;
+};
+export declare function createRedisPiiVaultStorage(redis: RedisPiiVaultClient, options?: RedisPiiVaultStorageOptions): PiiVaultStorage;
+export declare class InMemoryPiiVaultStorage implements PiiVaultStorage {
+    private scopes;
+    private tokenIndex;
+    get(scopeId: string, token: string): string | undefined;
+    getByToken(token: string): string | undefined;
+    set(scopeId: string, token: string, value: string): void;
+    entries(scopeId: string): PiiVaultEntry[];
+    private getScope;
+}
+export declare const defaultPiiVaultStorage: InMemoryPiiVaultStorage;

package/dist/pii/storage.js

@@ -0,0 +1,64 @@
+export function createPiiVaultStorage(adapter) {
+    return adapter;
+}
+export function createRedisPiiVaultStorage(redis, options = {}) {
+    const keyPrefix = options.keyPrefix ?? 'genkit-guard:pii';
+    const tokenIndexKey = options.tokenIndexKey ?? `${keyPrefix}:tokens`;
+    const hGet = redis.hGet?.bind(redis) ?? redis.hget?.bind(redis);
+    const hSet = redis.hSet?.bind(redis) ?? redis.hset?.bind(redis);
+    const hGetAll = redis.hGetAll?.bind(redis) ?? redis.hgetall?.bind(redis);
+    if (!hGet || !hSet || !hGetAll) {
+        throw new Error('Redis PII vault storage requires hGet/hSet/hGetAll or hget/hset/hgetall methods.');
+    }
+    const scopeKey = (scopeId) => `${keyPrefix}:scope:${scopeId}`;
+    async function maybeExpire(key) {
+        if (options.ttlSeconds && redis.expire) {
+            await redis.expire(key, options.ttlSeconds);
+        }
+    }
+    return createPiiVaultStorage({
+        async get(scopeId, token) {
+            return (await hGet(scopeKey(scopeId), token)) ?? undefined;
+        },
+        async getByToken(token) {
+            return (await hGet(tokenIndexKey, token)) ?? undefined;
+        },
+        async set(scopeId, token, value) {
+            const scopedKey = scopeKey(scopeId);
+            await hSet(scopedKey, token, value);
+            await hSet(tokenIndexKey, token, value);
+            await maybeExpire(scopedKey);
+            await maybeExpire(tokenIndexKey);
+        },
+        async entries(scopeId) {
+            const values = await hGetAll(scopeKey(scopeId));
+            return Object.entries(values).map(([token, value]) => ({ token, value }));
+        },
+    });
+}
+export class InMemoryPiiVaultStorage {
+    scopes = new Map();
+    tokenIndex = new Map();
+    get(scopeId, token) {
+        return this.scopes.get(scopeId)?.get(token);
+    }
+    getByToken(token) {
+        return this.tokenIndex.get(token);
+    }
+    set(scopeId, token, value) {
+        this.getScope(scopeId).set(token, value);
+        this.tokenIndex.set(token, value);
+    }
+    entries(scopeId) {
+        return Array.from(this.getScope(scopeId), ([token, value]) => ({ token, value }));
+    }
+    getScope(scopeId) {
+        let scope = this.scopes.get(scopeId);
+        if (!scope) {
+            scope = new Map();
+            this.scopes.set(scopeId, scope);
+        }
+        return scope;
+    }
+}
+export const defaultPiiVaultStorage = new InMemoryPiiVaultStorage();

package/dist/pii/tokenizer.d.ts

@@ -1,19 +1,29 @@
+import { type PiiVaultStorage } from './storage.js';
 export type PiiResult = {
     maskedText: string;
     pii: Record<string, string>;
     piiTypes: string[];
 };
+export type PiiTokenizerOptions = {
+    scopeId?: string;
+    storage?: PiiVaultStorage;
+};
 export declare class PiiTokenizer {
-    private vault;
+    private valueToToken;
     private counter;
     private piiTypes;
+    private scopeId;
+    private tokenNamespace;
+    private storage;
+    constructor(options?: PiiTokenizerOptions);
     private createToken;
     mask(text: string, matches: {
         type: string;
         value: string;
-    }[]): PiiResult;
-    unmask(text: string): string;
-    getVault(): {
+    }[]): Promise<PiiResult>;
+    unmask(text: string): Promise<string>;
+    importTokens(text: string): Promise<void>;
+    getVault(): Promise<{
         [k: string]: string;
-    };
+    }>;
 }

package/dist/pii/tokenizer.js

@@ -1,32 +1,69 @@
+import { defaultPiiVaultStorage } from './storage.js';
 export class PiiTokenizer {
-    vault = new Map();
+    valueToToken = new Map();
     counter = 0;
     piiTypes = new Set();
+    scopeId;
+    tokenNamespace;
+    storage;
+    constructor(options = {}) {
+        this.scopeId = options.scopeId ?? createVaultScopeId();
+        this.tokenNamespace = createVaultScopeId();
+        this.storage = options.storage ?? defaultPiiVaultStorage;
+    }
     createToken(type) {
-        return `[[${type}_${this.counter++}]]`;
+        return `[[${type}_${this.tokenNamespace}_${this.counter++}]]`;
     }
-    mask(text, matches) {
+    async mask(text, matches) {
         let masked = text;
         for (const match of matches) {
-            const token = this.createToken(match.type);
-            this.vault.set(token, match.value);
+            if (!masked.includes(match.value)) {
+                continue;
+            }
+            const key = `${match.type}:${match.value}`;
+            let token = this.valueToToken.get(key);
+            if (!token) {
+                token = this.createToken(match.type);
+                this.valueToToken.set(key, token);
+                await this.storage.set(this.scopeId, token, match.value);
+            }
             this.piiTypes.add(match.type.toLowerCase());
             masked = masked.split(match.value).join(token);
         }
         return {
             maskedText: masked,
-            pii: Object.fromEntries(this.vault),
+            pii: await this.getVault(),
             piiTypes: Array.from(this.piiTypes)
         };
     }
-    unmask(text) {
+    async unmask(text) {
         let result = text;
-        this.vault.forEach((value, token) => {
-            result = result.split(token).join(value); // Global replacement
-        });
+        const entries = await this.storage.entries(this.scopeId);
+        for (const { token, value } of entries) {
+            result = result.split(token).join(value);
+        }
         return result;
     }
-    getVault() {
-        return Object.fromEntries(this.vault);
+    async importTokens(text) {
+        if (!this.storage.getByToken) {
+            return;
+        }
+        const tokenMatches = Array.from(text.matchAll(/\[\[[A-Z_]+_[A-Za-z0-9]+_\d+\]\]/g));
+        for (const [token] of tokenMatches) {
+            const value = await this.storage.getByToken(token);
+            if (value) {
+                await this.storage.set(this.scopeId, token, value);
+            }
+        }
+    }
+    async getVault() {
+        const entries = await this.storage.entries(this.scopeId);
+        return Object.fromEntries(entries.map(({ token, value }) => [token, value]));
+    }
+}
+function createVaultScopeId() {
+    if (typeof crypto !== 'undefined' && 'randomUUID' in crypto) {
+        return crypto.randomUUID().replace(/-/g, '');
     }
+    return `${Date.now().toString(36)}${Math.random().toString(36).slice(2)}`;
 }

package/dist/util/singleton.js

@@ -21,7 +21,22 @@ export class ModelSingleton {
         env.localModelPath = modelPath;
         // Allow remote download if missing
         env.allowRemoteModels = true;
-        console.log("[Guard] Using model directory:", modelPath);
+        console.log(JSON.stringify({
+            timestamp: new Date().toISOString(),
+            severityText: 'INFO',
+            severityNumber: 9,
+            body: 'Using guard model directory',
+            resource: {
+                attributes: {
+                    'service.name': '@intflows/genkit-guard',
+                },
+            },
+            attributes: {
+                'event.name': 'guard.models.directory',
+                'code.namespace': 'genkit-guard',
+                modelPath,
+            },
+        }));
     }
     static async getExtractor(modelName = 'Xenova/all-MiniLM-L6-v2') {
         if (!this.extractors.has(modelName)) {

package/package.json

@@ -22,7 +22,7 @@
     "huggingface"
   ],
   "license": "Apache-2.0",
-  "version": "0.0.9",
+  "version": "0.0.11",
   "type": "module",
   "exports": "./dist/index.js",
   "types": "./dist/index.d.ts",
@@ -35,6 +35,11 @@
   "scripts": {
     "prepare-models": "node scripts/download-model.js",
     "build": "tsc",
+    "test": "npm run test:types && npm run test:storage && npm run test:concurrency",
+    "test:types": "tsc --noEmit --ignoreConfig --module NodeNext --moduleResolution NodeNext --target ESNext --strict --skipLibCheck scripts/test-types.ts",
+    "test:storage": "npm run build && node scripts/test-storage.js",
+    "test:concurrency": "npm run build && node scripts/test-concurrent-vault.js",
+    "test:redis": "npm run build && node scripts/test-redis-vault.js",
     "prepublishOnly": "npm run build"
   },
   "dependencies": {
@@ -42,11 +47,11 @@
     "zod": "^4.4.3"
   },
   "peerDependencies": {
-    "genkit": "^1.37.0"
+    "genkit": "1.39.0"
   },
   "devDependencies": {
     "@types/node": "^25.8.0",
-    "genkit": "^1.37.0",
+    "genkit": "^1.39.0",
     "typescript": "^6.0.3"
   }
 }