@intflows/genkit-guard 0.0.9-alpha.1 → 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, 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,6 +1,7 @@
 import { ModelSingleton } from './util/singleton.js';
 // export { intentGuard, piiGuard  } 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({

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,32 @@ 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>;
@@ -94,14 +113,32 @@ 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>;
@@ -156,14 +193,32 @@ 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>;
@@ -186,7 +241,14 @@ declare const guardConfigSchema: z.ZodObject<{
         extractor?: string | undefined;
     }>>;
 }, z.ZodTypeAny, "passthrough">>;
-type GuardConfig = z.infer<typeof guardConfigSchema>;
+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>;
@@ -220,14 +282,32 @@ 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>;
@@ -282,14 +362,32 @@ 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>;
@@ -344,14 +442,32 @@ 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>;

package/dist/middleware/middleware.js

@@ -2,6 +2,7 @@ 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({
@@ -16,6 +17,10 @@ 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(),
@@ -90,9 +95,10 @@ function createGuardHooks(config) {
             const textForPii = collectModelRequestText(req);
             const piiResponse = await scanPII(textForPii, config);
             const piiMatches = piiResponse?.matches || [];
-            const tokenizer = new PiiTokenizer();
+            const tokenizer = createTokenizer(config, req, ctx);
             const piiTypes = uniqueTypes(piiMatches);
-            maskModelRequest(req, tokenizer, 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,
@@ -113,7 +119,7 @@ function createGuardHooks(config) {
                 piiClassifierOutput: piiResponse.classifier,
             };
             const res = await next(req, ctx);
-            const unmaskedResponse = unmaskObject(res, [tokenizer]);
+            const unmaskedResponse = await unmaskObject(res, getGuardState(ctx).tokenizers);
             logger('info', 'guard.model.response.unmasked', 'Model response unmasked for downstream execution', {
                 piiTypes,
             });
@@ -123,7 +129,7 @@ function createGuardHooks(config) {
             const state = getGuardState(ctx);
             const toolName = req?.toolRequest?.name;
             if (req?.toolRequest && 'input' in req.toolRequest) {
-                req.toolRequest.input = unmaskObject(req.toolRequest.input, state.tokenizers);
+                req.toolRequest.input = await unmaskObject(req.toolRequest.input, state.tokenizers);
             }
             const toolInputText = collectStrings(req?.toolRequest?.input).join('\n');
             const piiResponse = await scanPII(toolInputText, config);
@@ -178,39 +184,39 @@ function collectModelRequestText(req) {
         ...collectStrings(req?.docs),
     ].join('\n');
 }
-function maskModelRequest(req, tokenizer, matches) {
+async function maskModelRequest(req, tokenizer, matches) {
     if (typeof req.prompt === 'string') {
-        req.prompt = tokenizer.mask(req.prompt, matches).maskedText;
+        req.prompt = (await tokenizer.mask(req.prompt, matches)).maskedText;
     }
     if (req.messages) {
-        req.messages = transformStrings(req.messages, (value) => tokenizer.mask(value, matches).maskedText);
+        req.messages = await transformStrings(req.messages, async (value) => (await tokenizer.mask(value, matches)).maskedText);
     }
     if (req.docs) {
-        req.docs = transformStrings(req.docs, (value) => tokenizer.mask(value, matches).maskedText);
+        req.docs = await transformStrings(req.docs, async (value) => (await tokenizer.mask(value, matches)).maskedText);
     }
 }
-function unmaskObject(obj, tokenizers) {
-    return transformStrings(obj, (value) => {
+async function unmaskObject(obj, tokenizers) {
+    return transformStrings(obj, async (value) => {
         let result = value;
         for (const tokenizer of tokenizers) {
-            result = tokenizer.unmask(result);
+            result = await tokenizer.unmask(result);
         }
         return result;
     });
 }
-function transformStrings(obj, transform) {
+async function transformStrings(obj, transform) {
     if (typeof obj === 'string') {
-        return transform(obj);
+        return await transform(obj);
     }
     if (Array.isArray(obj)) {
         for (let i = 0; i < obj.length; i++) {
-            obj[i] = transformStrings(obj[i], transform);
+            obj[i] = await transformStrings(obj[i], transform);
         }
         return obj;
     }
     if (obj !== null && typeof obj === 'object') {
         for (const key of Object.keys(obj)) {
-            obj[key] = transformStrings(obj[key], transform);
+            obj[key] = await transformStrings(obj[key], transform);
         }
         return obj;
     }
@@ -249,6 +255,16 @@ 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())));
 }

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,20 +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,12 +1,20 @@
+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) {
             if (!masked.includes(match.value)) {
@@ -17,25 +25,45 @@ export class PiiTokenizer {
             if (!token) {
                 token = this.createToken(match.type);
                 this.valueToToken.set(key, token);
-                this.vault.set(token, match.value);
+                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/package.json

@@ -22,7 +22,7 @@
     "huggingface"
   ],
   "license": "Apache-2.0",
-  "version": "0.0.9-alpha.1",
+  "version": "0.0.10",
   "type": "module",
   "exports": "./dist/index.js",
   "types": "./dist/index.d.ts",
@@ -32,11 +32,12 @@
     "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"

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.`);