@intflows/genkit-guard 0.0.9 → 0.0.10

package/README.md

@@ -1,7 +1,5 @@
 # **@intflows/genkit-guard** 
 
-#### _This version uses OpenAI/privacy-filter instead of bert-base-NER_
-
 ### **Lightweight Intent, PII, and Safety Guardrails for Genkit**
 
 `@intflows/genkit-guard` provides a modular guardrail layer for Genkit flows.  
@@ -114,6 +112,29 @@ const response = await ai.generate({
 
 ![Image showing PII data masked ](./MaskedPII.png)
 
+---
+## Example
+
+An example genkit flow is present in `example` directory.
+
+```bash
+git clone https://github.com/IntFlows/genkit-guard.git
+cd genkit-guard/example
+npm install
+node node_modules/@intflows/genkit-guard/scripts/download-model.js
+npx tsx src/index.ts
+```
+
+Or you can run the flow with genkit dev UI
+
+```bash
+git clone https://github.com/IntFlows/genkit-guard.git
+cd genkit-guard/example
+npm install
+node node_modules/@intflows/genkit-guard/scripts/download-model.js
+genkit start -- npx tsx src/index.ts
+```
+
 ---
 
 ## 🧠 How It Works
@@ -177,6 +198,51 @@ pii: {
 }
 ```
 
+### **PII Vault Isolation and External Storage**
+
+By default, PII is stored in an in-memory vault scoped to a single tokenizer instance. Tokens include a generated vault scope:
+
+```txt
+"Email john.doe@example.com" -> "Email [[EMAIL_<namespace>_0]]"
+```
+
+That generated namespace prevents two concurrent calls from sharing the same visible placeholder names. Vault lookups are isolated by the configured storage scope, so User A and User B can safely produce their own email tokens without cross-resolving each other's PII.
+
+For applications that need persistence, distributed workers, audits, or tenant-specific storage, provide a vault storage backend:
+
+```ts
+import { guard, type PiiVaultStorage } from "@intflows/genkit-guard";
+
+class RedisPiiVaultStorage implements PiiVaultStorage {
+  constructor(private redis: RedisClient) {}
+
+  async get(scopeId: string, token: string) {
+    return this.redis.hget(`pii:${scopeId}`, token);
+  }
+
+  async set(scopeId: string, token: string, value: string) {
+    await this.redis.hset(`pii:${scopeId}`, token, value);
+  }
+
+  async entries(scopeId: string) {
+    const values = await this.redis.hgetall(`pii:${scopeId}`);
+    return Object.entries(values).map(([token, value]) => ({ token, value }));
+  }
+}
+
+guard({
+  pii: {
+    reversible: true,
+    vault: {
+      storage: new RedisPiiVaultStorage(redis),
+      scopeId: (req, ctx) => ctx?.auth?.sessionId ?? req?.metadata?.requestId
+    }
+  }
+});
+```
+
+Choose a `scopeId` that matches your isolation boundary, such as request ID, session ID, tenant/user ID, or a combination like `tenantId:userId:requestId`. A shared external backend should never ignore `scopeId`, because placeholders are only safe when resolved against the correct vault scope. The placeholder sent to the model uses an opaque generated namespace rather than exposing your `scopeId`.
+
 ---
 
 ## 🛡️ Why This Library Exists

package/dist/index.d.ts

@@ -1,4 +1,6 @@
-export { guard } from './middleware/middleware.js';
+export { guard, guardAction, guardMiddleware, guardPlugin } from './middleware/middleware.js';
+export { InMemoryPiiVaultStorage, defaultPiiVaultStorage } from './pii/storage.js';
+export type { PiiVaultEntry, PiiVaultStorage } from './pii/storage.js';
 export * from './core/types.js';
 /**
  * Pre-load the model to avoid cold-start delay on first user request.

package/dist/index.js

@@ -1,12 +1,31 @@
 import { ModelSingleton } from './util/singleton.js';
 // export { intentGuard, piiGuard  } from './middleware/middleware.js';
-export { guard } from './middleware/middleware.js';
+export { guard, guardAction, guardMiddleware, guardPlugin } from './middleware/middleware.js';
+export { InMemoryPiiVaultStorage, defaultPiiVaultStorage } from './pii/storage.js';
 export * from './core/types.js';
+function logGuardEvent(eventName, body, attributes = {}) {
+    console.log(JSON.stringify({
+        timestamp: new Date().toISOString(),
+        severityText: 'INFO',
+        severityNumber: 9,
+        body,
+        resource: {
+            attributes: {
+                'service.name': '@intflows/genkit-guard',
+            },
+        },
+        attributes: {
+            'event.name': eventName,
+            'code.namespace': 'genkit-guard',
+            ...attributes,
+        },
+    }));
+}
 /**
  * Pre-load the model to avoid cold-start delay on first user request.
  */
 export async function initGuard(config) {
-    console.log('[Guard] Loading local models...');
+    logGuardEvent('guard.models.loading', 'Loading local guard models');
     const extractorModel = config?.models?.extractor ?? 'Xenova/all-MiniLM-L6-v2';
     const piiModel = config?.pii?.model;
     const piiMode = config?.pii?.mode ?? 'ner';
@@ -18,5 +37,7 @@ export async function initGuard(config) {
         tasks.push(ModelSingleton.getPIIClassifier(piiModel ?? 'openai/privacy-filter'));
     }
     await Promise.all(tasks);
-    console.log('[Guard] Models loaded');
+    logGuardEvent('guard.models.loaded', 'Local guard models loaded', {
+        piiMode,
+    });
 }

package/dist/middleware/middleware.d.ts

@@ -1,4 +1,5 @@
 import { z } from 'genkit';
+import { type PiiVaultStorage } from '../pii/storage.js';
 declare const guardConfigSchema: z.ZodObject<{
     intent: z.ZodOptional<z.ZodObject<{
         mode: z.ZodOptional<z.ZodString>;
@@ -32,14 +33,45 @@ declare const guardConfigSchema: z.ZodObject<{
         reversible: z.ZodOptional<z.ZodBoolean>;
         model: z.ZodOptional<z.ZodString>;
         mode: z.ZodOptional<z.ZodEnum<["ner", "classifier"]>>;
+        vault: z.ZodOptional<z.ZodObject<{
+            storage: z.ZodOptional<z.ZodAny>;
+            scopeId: z.ZodOptional<z.ZodAny>;
+        }, "strip", z.ZodTypeAny, {
+            storage?: any;
+            scopeId?: any;
+        }, {
+            storage?: any;
+            scopeId?: any;
+        }>>;
     }, "strip", z.ZodTypeAny, {
         mode?: "ner" | "classifier" | undefined;
         reversible?: boolean | undefined;
         model?: string | undefined;
+        vault?: {
+            storage?: any;
+            scopeId?: any;
+        } | undefined;
     }, {
         mode?: "ner" | "classifier" | undefined;
         reversible?: boolean | undefined;
         model?: string | undefined;
+        vault?: {
+            storage?: any;
+            scopeId?: any;
+        } | undefined;
+    }>>;
+    logging: z.ZodOptional<z.ZodObject<{
+        enabled: z.ZodOptional<z.ZodBoolean>;
+        level: z.ZodOptional<z.ZodEnum<["debug", "info", "warn", "error"]>>;
+        serviceName: z.ZodOptional<z.ZodString>;
+    }, "strip", z.ZodTypeAny, {
+        enabled?: boolean | undefined;
+        level?: "debug" | "info" | "warn" | "error" | undefined;
+        serviceName?: string | undefined;
+    }, {
+        enabled?: boolean | undefined;
+        level?: "debug" | "info" | "warn" | "error" | undefined;
+        serviceName?: string | undefined;
     }>>;
     models: z.ZodOptional<z.ZodObject<{
         extractor: z.ZodOptional<z.ZodString>;
@@ -81,14 +113,45 @@ declare const guardConfigSchema: z.ZodObject<{
         reversible: z.ZodOptional<z.ZodBoolean>;
         model: z.ZodOptional<z.ZodString>;
         mode: z.ZodOptional<z.ZodEnum<["ner", "classifier"]>>;
+        vault: z.ZodOptional<z.ZodObject<{
+            storage: z.ZodOptional<z.ZodAny>;
+            scopeId: z.ZodOptional<z.ZodAny>;
+        }, "strip", z.ZodTypeAny, {
+            storage?: any;
+            scopeId?: any;
+        }, {
+            storage?: any;
+            scopeId?: any;
+        }>>;
     }, "strip", z.ZodTypeAny, {
         mode?: "ner" | "classifier" | undefined;
         reversible?: boolean | undefined;
         model?: string | undefined;
+        vault?: {
+            storage?: any;
+            scopeId?: any;
+        } | undefined;
     }, {
         mode?: "ner" | "classifier" | undefined;
         reversible?: boolean | undefined;
         model?: string | undefined;
+        vault?: {
+            storage?: any;
+            scopeId?: any;
+        } | undefined;
+    }>>;
+    logging: z.ZodOptional<z.ZodObject<{
+        enabled: z.ZodOptional<z.ZodBoolean>;
+        level: z.ZodOptional<z.ZodEnum<["debug", "info", "warn", "error"]>>;
+        serviceName: z.ZodOptional<z.ZodString>;
+    }, "strip", z.ZodTypeAny, {
+        enabled?: boolean | undefined;
+        level?: "debug" | "info" | "warn" | "error" | undefined;
+        serviceName?: string | undefined;
+    }, {
+        enabled?: boolean | undefined;
+        level?: "debug" | "info" | "warn" | "error" | undefined;
+        serviceName?: string | undefined;
     }>>;
     models: z.ZodOptional<z.ZodObject<{
         extractor: z.ZodOptional<z.ZodString>;
@@ -130,14 +193,45 @@ declare const guardConfigSchema: z.ZodObject<{
         reversible: z.ZodOptional<z.ZodBoolean>;
         model: z.ZodOptional<z.ZodString>;
         mode: z.ZodOptional<z.ZodEnum<["ner", "classifier"]>>;
+        vault: z.ZodOptional<z.ZodObject<{
+            storage: z.ZodOptional<z.ZodAny>;
+            scopeId: z.ZodOptional<z.ZodAny>;
+        }, "strip", z.ZodTypeAny, {
+            storage?: any;
+            scopeId?: any;
+        }, {
+            storage?: any;
+            scopeId?: any;
+        }>>;
     }, "strip", z.ZodTypeAny, {
         mode?: "ner" | "classifier" | undefined;
         reversible?: boolean | undefined;
         model?: string | undefined;
+        vault?: {
+            storage?: any;
+            scopeId?: any;
+        } | undefined;
     }, {
         mode?: "ner" | "classifier" | undefined;
         reversible?: boolean | undefined;
         model?: string | undefined;
+        vault?: {
+            storage?: any;
+            scopeId?: any;
+        } | undefined;
+    }>>;
+    logging: z.ZodOptional<z.ZodObject<{
+        enabled: z.ZodOptional<z.ZodBoolean>;
+        level: z.ZodOptional<z.ZodEnum<["debug", "info", "warn", "error"]>>;
+        serviceName: z.ZodOptional<z.ZodString>;
+    }, "strip", z.ZodTypeAny, {
+        enabled?: boolean | undefined;
+        level?: "debug" | "info" | "warn" | "error" | undefined;
+        serviceName?: string | undefined;
+    }, {
+        enabled?: boolean | undefined;
+        level?: "debug" | "info" | "warn" | "error" | undefined;
+        serviceName?: string | undefined;
     }>>;
     models: z.ZodOptional<z.ZodObject<{
         extractor: z.ZodOptional<z.ZodString>;
@@ -147,6 +241,14 @@ declare const guardConfigSchema: z.ZodObject<{
         extractor?: string | undefined;
     }>>;
 }, z.ZodTypeAny, "passthrough">>;
+type GuardConfig = z.infer<typeof guardConfigSchema> & {
+    pii?: {
+        vault?: {
+            storage?: PiiVaultStorage;
+            scopeId?: string | ((req: any, ctx: any) => string | undefined);
+        };
+    };
+};
 export declare const guardMiddleware: import("genkit").GenerateMiddleware<z.ZodObject<{
     intent: z.ZodOptional<z.ZodObject<{
         mode: z.ZodOptional<z.ZodString>;
@@ -180,14 +282,45 @@ export declare const guardMiddleware: import("genkit").GenerateMiddleware<z.ZodO
         reversible: z.ZodOptional<z.ZodBoolean>;
         model: z.ZodOptional<z.ZodString>;
         mode: z.ZodOptional<z.ZodEnum<["ner", "classifier"]>>;
+        vault: z.ZodOptional<z.ZodObject<{
+            storage: z.ZodOptional<z.ZodAny>;
+            scopeId: z.ZodOptional<z.ZodAny>;
+        }, "strip", z.ZodTypeAny, {
+            storage?: any;
+            scopeId?: any;
+        }, {
+            storage?: any;
+            scopeId?: any;
+        }>>;
     }, "strip", z.ZodTypeAny, {
         mode?: "ner" | "classifier" | undefined;
         reversible?: boolean | undefined;
         model?: string | undefined;
+        vault?: {
+            storage?: any;
+            scopeId?: any;
+        } | undefined;
     }, {
         mode?: "ner" | "classifier" | undefined;
         reversible?: boolean | undefined;
         model?: string | undefined;
+        vault?: {
+            storage?: any;
+            scopeId?: any;
+        } | undefined;
+    }>>;
+    logging: z.ZodOptional<z.ZodObject<{
+        enabled: z.ZodOptional<z.ZodBoolean>;
+        level: z.ZodOptional<z.ZodEnum<["debug", "info", "warn", "error"]>>;
+        serviceName: z.ZodOptional<z.ZodString>;
+    }, "strip", z.ZodTypeAny, {
+        enabled?: boolean | undefined;
+        level?: "debug" | "info" | "warn" | "error" | undefined;
+        serviceName?: string | undefined;
+    }, {
+        enabled?: boolean | undefined;
+        level?: "debug" | "info" | "warn" | "error" | undefined;
+        serviceName?: string | undefined;
     }>>;
     models: z.ZodOptional<z.ZodObject<{
         extractor: z.ZodOptional<z.ZodString>;
@@ -229,14 +362,45 @@ export declare const guardMiddleware: import("genkit").GenerateMiddleware<z.ZodO
         reversible: z.ZodOptional<z.ZodBoolean>;
         model: z.ZodOptional<z.ZodString>;
         mode: z.ZodOptional<z.ZodEnum<["ner", "classifier"]>>;
+        vault: z.ZodOptional<z.ZodObject<{
+            storage: z.ZodOptional<z.ZodAny>;
+            scopeId: z.ZodOptional<z.ZodAny>;
+        }, "strip", z.ZodTypeAny, {
+            storage?: any;
+            scopeId?: any;
+        }, {
+            storage?: any;
+            scopeId?: any;
+        }>>;
     }, "strip", z.ZodTypeAny, {
         mode?: "ner" | "classifier" | undefined;
         reversible?: boolean | undefined;
         model?: string | undefined;
+        vault?: {
+            storage?: any;
+            scopeId?: any;
+        } | undefined;
     }, {
         mode?: "ner" | "classifier" | undefined;
         reversible?: boolean | undefined;
         model?: string | undefined;
+        vault?: {
+            storage?: any;
+            scopeId?: any;
+        } | undefined;
+    }>>;
+    logging: z.ZodOptional<z.ZodObject<{
+        enabled: z.ZodOptional<z.ZodBoolean>;
+        level: z.ZodOptional<z.ZodEnum<["debug", "info", "warn", "error"]>>;
+        serviceName: z.ZodOptional<z.ZodString>;
+    }, "strip", z.ZodTypeAny, {
+        enabled?: boolean | undefined;
+        level?: "debug" | "info" | "warn" | "error" | undefined;
+        serviceName?: string | undefined;
+    }, {
+        enabled?: boolean | undefined;
+        level?: "debug" | "info" | "warn" | "error" | undefined;
+        serviceName?: string | undefined;
     }>>;
     models: z.ZodOptional<z.ZodObject<{
         extractor: z.ZodOptional<z.ZodString>;
@@ -278,14 +442,45 @@ export declare const guardMiddleware: import("genkit").GenerateMiddleware<z.ZodO
         reversible: z.ZodOptional<z.ZodBoolean>;
         model: z.ZodOptional<z.ZodString>;
         mode: z.ZodOptional<z.ZodEnum<["ner", "classifier"]>>;
+        vault: z.ZodOptional<z.ZodObject<{
+            storage: z.ZodOptional<z.ZodAny>;
+            scopeId: z.ZodOptional<z.ZodAny>;
+        }, "strip", z.ZodTypeAny, {
+            storage?: any;
+            scopeId?: any;
+        }, {
+            storage?: any;
+            scopeId?: any;
+        }>>;
     }, "strip", z.ZodTypeAny, {
         mode?: "ner" | "classifier" | undefined;
         reversible?: boolean | undefined;
         model?: string | undefined;
+        vault?: {
+            storage?: any;
+            scopeId?: any;
+        } | undefined;
     }, {
         mode?: "ner" | "classifier" | undefined;
         reversible?: boolean | undefined;
         model?: string | undefined;
+        vault?: {
+            storage?: any;
+            scopeId?: any;
+        } | undefined;
+    }>>;
+    logging: z.ZodOptional<z.ZodObject<{
+        enabled: z.ZodOptional<z.ZodBoolean>;
+        level: z.ZodOptional<z.ZodEnum<["debug", "info", "warn", "error"]>>;
+        serviceName: z.ZodOptional<z.ZodString>;
+    }, "strip", z.ZodTypeAny, {
+        enabled?: boolean | undefined;
+        level?: "debug" | "info" | "warn" | "error" | undefined;
+        serviceName?: string | undefined;
+    }, {
+        enabled?: boolean | undefined;
+        level?: "debug" | "info" | "warn" | "error" | undefined;
+        serviceName?: string | undefined;
     }>>;
     models: z.ZodOptional<z.ZodObject<{
         extractor: z.ZodOptional<z.ZodString>;
@@ -296,5 +491,6 @@ export declare const guardMiddleware: import("genkit").GenerateMiddleware<z.ZodO
     }>>;
 }, z.ZodTypeAny, "passthrough">>, void>;
 export declare const guardPlugin: (pluginOptions: void) => import("@genkit-ai/ai").GenkitPluginV2;
-export declare function guard(config?: z.infer<typeof guardConfigSchema>): (req: any, ctxOrNext: any, maybeNext?: any) => Promise<any>;
+export declare function guard(config?: GuardConfig): (req: any, ctxOrNext: any, maybeNext?: any) => Promise<any>;
+export declare const guardAction: typeof guard;
 export {};

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,20 @@
+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 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,26 @@
+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.10",
   "type": "module",
   "exports": "./dist/index.js",
   "types": "./dist/index.d.ts",
@@ -32,21 +32,22 @@
     "scripts/",
     "dist/"
   ],
-  "scripts": {
-    "prepare-models": "node scripts/download-model.js",
-    "build": "tsc",
-    "prepublishOnly": "npm run build"
-  },
+  "scripts": {
+    "prepare-models": "node scripts/download-model.js",
+    "build": "tsc",
+    "test:concurrency": "npm run build && node scripts/test-concurrent-vault.js",
+    "prepublishOnly": "npm run build"
+  },
   "dependencies": {
     "@huggingface/transformers": "^4.2.0",
     "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"
   }
 }

package/scripts/test-concurrent-vault.js

@@ -0,0 +1,96 @@
+import assert from 'node:assert/strict';
+import { InMemoryPiiVaultStorage } from '../dist/index.js';
+import { PiiTokenizer } from '../dist/pii/tokenizer.js';
+
+const storage = new InMemoryPiiVaultStorage();
+const userCount = 100;
+
+const users = Array.from({ length: userCount }, (_, index) => ({
+  scopeId: `tenantA:user${index}`,
+  email: `user${index}@example.com`,
+}));
+
+const results = await Promise.all(users.map(async ({ scopeId, email }) => {
+  const tokenizer = new PiiTokenizer({ scopeId, storage });
+  const prompt = `Fetch metadata for the blob named 'sensitive-user-data-${email}.json' inside Azure Blob Storage.`;
+  const masked = await tokenizer.mask(prompt, [{ type: 'EMAIL', value: email }]);
+
+  const simulatedModelToolInput = {
+    blobName: masked.maskedText.match(/'([^']+)'/)?.[1],
+  };
+
+  const toolInput = {
+    blobName: await tokenizer.unmask(simulatedModelToolInput.blobName),
+  };
+
+  const simulatedModelResponse = `The blob belongs to ${masked.maskedText.match(/\[\[EMAIL_[^\]]+\]\]/)?.[0]}.`;
+  const response = await tokenizer.unmask(simulatedModelResponse);
+
+  return {
+    scopeId,
+    email,
+    prompt,
+    maskedText: masked.maskedText,
+    toolInput,
+    response,
+  };
+}));
+
+const placeholders = new Set(results.map((result) => result.maskedText));
+assert.equal(placeholders.size, userCount, 'each concurrent user should get a unique scoped placeholder');
+
+for (const result of results) {
+  assert.equal(result.toolInput.blobName, `sensitive-user-data-${result.email}.json`);
+  assert.equal(result.response, `The blob belongs to ${result.email}.`);
+}
+
+const first = results[0];
+const second = results[1];
+const firstTokenizer = new PiiTokenizer({ scopeId: first.scopeId, storage });
+assert.equal(
+  await firstTokenizer.unmask(second.maskedText),
+  second.maskedText,
+  'one user scope must not resolve another user scope placeholder'
+);
+
+const promptTokenizer = new PiiTokenizer({ scopeId: 'trace:request', storage });
+const promptMasked = await promptTokenizer.mask(
+  "Fetch metadata for the blob named 'sensitive-user-data-john.doe@example.com.json' inside Azure Blob Storage.",
+  [{ type: 'EMAIL', value: 'john.doe@example.com' }]
+);
+
+const toolResponseTokenizer = new PiiTokenizer({ scopeId: 'trace:request', storage });
+await toolResponseTokenizer.mask(
+  'john.doe@example.com',
+  [{ type: 'EMAIL', value: 'john.doe@example.com' }]
+);
+
+const finalModelResponse = `The metadata for the blob '${promptMasked.maskedText.match(/\[\[EMAIL_[^\]]+\]\]/)?.[0]}' has been successfully retrieved.`;
+let unmaskedFinalResponse = finalModelResponse;
+for (const tokenizer of [promptTokenizer, toolResponseTokenizer]) {
+  unmaskedFinalResponse = await tokenizer.unmask(unmaskedFinalResponse);
+}
+
+assert.equal(
+  unmaskedFinalResponse,
+  "The metadata for the blob 'john.doe@example.com' has been successfully retrieved.",
+  'final model output should unmask placeholders created by earlier middleware passes'
+);
+
+const firstMiddlewarePass = new PiiTokenizer({ storage });
+const firstPassMasked = await firstMiddlewarePass.mask(
+  "Fetch metadata for the blob named 'sensitive-user-data-jane.doe@example.com.json' inside Azure Blob Storage.",
+  [{ type: 'EMAIL', value: 'sensitive-user-data-jane.doe@example.com.json' }]
+);
+const laterMiddlewarePass = new PiiTokenizer({ storage });
+const crossContextResponse = `The metadata for the blob '${firstPassMasked.maskedText.match(/\[\[EMAIL_[^\]]+\]\]/)?.[0]}' has been successfully retrieved.`;
+
+await laterMiddlewarePass.importTokens(firstPassMasked.maskedText);
+
+assert.equal(
+  await laterMiddlewarePass.unmask(crossContextResponse),
+  "The metadata for the blob 'sensitive-user-data-jane.doe@example.com.json' has been successfully retrieved.",
+  'a later middleware context should resolve an earlier opaque token from the shared default vault'
+);
+
+console.log(`Concurrent PII vault simulation passed for ${userCount} isolated scopes.`);