@intflows/genkit-guard 0.0.11 → 0.0.13

package/README.md

@@ -63,7 +63,7 @@ npm install @intflows/genkit-guard
 # Download Local Models (Only needed once)
 node node_modules/@intflows/genkit-guard/scripts/download-model.js
 ```
-_ This downloads the models to ./models folder, the total size is ~1.5 GB ( 1GB for Openai/privacy-filter + .5 GB for MiniLM-L6-v2) 
+_This downloads the models to `./models`; the total size is approximately 1.5 GB._
 
 ### 2. Update genkit
 
@@ -193,10 +193,26 @@ intent: {
 ### **PII Guard**
 
 ```ts
-pii: {
-  reversible: true
-}
-```
+pii: {
+  reversible: true,
+  mode: "classifier"
+}
+```
+
+`classifier` mode uses `openai/privacy-filter` as a token-classification model with aggregated
+spans. Model-detected names, addresses, emails, phone numbers, URLs, dates, account numbers and
+secrets are converted into reversible masking tokens. Regex rules continue to run as an additional
+layer, and duplicate spans are masked only once.
+
+During multi-turn tool execution, opaque tokens returned through a different Genkit middleware
+context are rehydrated from the configured vault before tool execution and before the final
+response is returned to the application.
+
+Preload the same mode during application startup:
+
+```ts
+await initGuard({ pii: { mode: "classifier" } });
+```
 
 ### **PII Vault Isolation and External Storage**
 
@@ -208,6 +224,7 @@ By default, PII is stored in an in-memory vault scoped to a single tokenizer ins
 
 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. Redis clients can be passed through the built-in helper:
 For applications that need persistence, distributed workers, audits, or tenant-specific storage, provide a vault storage backend. Redis clients can be passed through the built-in helper:
 
 ```ts
@@ -223,7 +240,8 @@ guard({
     vault: {
       storage: createRedisPiiVaultStorage(redis, {
         keyPrefix: "my-app:pii",
-        ttlSeconds: 3600
+        ttlSeconds: 3600,
+        fallbackToMemory: true
       }),
       scopeId: (req, ctx) => ctx?.auth?.sessionId ?? req?.metadata?.requestId
     }
@@ -231,6 +249,12 @@ guard({
 });
 ```
 
+`ttlSeconds` applies the configured expiry to both the scoped vault and token index. When
+`fallbackToMemory` is enabled, successful writes are also mirrored in process memory and Redis
+operation failures fall back to that mirror. The fallback is disabled by default, is local to one
+process, and is not a replacement for Redis persistence or multi-worker availability. Its in-memory
+entries observe the same TTL. Redis errors continue to propagate when fallback is disabled.
+
 For another backend, use `createPiiVaultStorage({ get, set, entries, getByToken })` with your database, cache, or secret store.
 
 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`.

package/dist/middleware/middleware.js

@@ -128,6 +128,11 @@ function createGuardHooks(config) {
         tool: async (req, ctx, next) => {
             const state = getGuardState(ctx);
             const toolName = req?.toolRequest?.name;
+            // Genkit may provide a fresh middleware context for a tool turn. Create a recovery
+            // tokenizer that uses the configured vault so opaque tokens can be rehydrated safely.
+            if (state.tokenizers.length === 0) {
+                state.tokenizers.push(createTokenizer(config, req, ctx));
+            }
             if (req?.toolRequest && 'input' in req.toolRequest) {
                 req.toolRequest.input = await unmaskObject(req.toolRequest.input, state.tokenizers);
             }
@@ -199,6 +204,9 @@ async function unmaskObject(obj, tokenizers) {
     return transformStrings(obj, async (value) => {
         let result = value;
         for (const tokenizer of tokenizers) {
+            // A token may have been produced by another model/tool turn with a different Genkit
+            // context or vault scope. Import only opaque tokens actually present in this value.
+            await tokenizer.importTokens(result);
             result = await tokenizer.unmask(result);
         }
         return result;

package/dist/pii/detector.d.ts

@@ -1,10 +1,20 @@
+export type PiiMatch = {
+    type: string;
+    value: string;
+};
+export type PrivacyFilterSpan = {
+    entity_group?: string;
+    entity?: string;
+    word?: string;
+    start?: number;
+    end?: number;
+    score?: number;
+};
 export declare function detectPII(text: string, opts?: {
     model?: string;
     mode?: 'ner' | 'classifier';
 }): Promise<{
-    matches: {
-        type: string;
-        value: string;
-    }[];
+    matches: PiiMatch[];
     classifier: any;
 }>;
+export declare function privacyFilterOutputToMatches(text: string, output: unknown): PiiMatch[];

package/dist/pii/detector.js

@@ -1,4 +1,14 @@
 import { ModelSingleton } from '../util/singleton.js';
+const PRIVACY_FILTER_TYPE_MAP = {
+    account_number: 'ACCOUNT_NUMBER',
+    private_address: 'ADDRESS',
+    private_email: 'EMAIL',
+    private_person: 'NAME',
+    private_phone: 'PHONE',
+    private_url: 'URL',
+    private_date: 'DATE',
+    secret: 'SECRET',
+};
 const REGEX_RULES = [
     // EMAIL (keep your existing one)
     { type: 'EMAIL', pattern: /\b[\w\.-]+@[\w\.-]+\.\w{2,}\b/gi },
@@ -36,12 +46,52 @@ export async function detectPII(text, opts) {
         }
     }
     else {
-        // classifier mode: we call the classifier and return its output alongside regex matches.
+        // Privacy Filter is a token-classification model. Aggregation produces complete spans
+        // rather than individual BIOES-labelled tokens.
         const cls = await ModelSingleton.getPIIClassifier(model);
-        classifierOutput = await cls(text);
+        classifierOutput = await cls(text, { aggregation_strategy: 'simple' });
+        for (const match of privacyFilterOutputToMatches(text, classifierOutput)) {
+            if (!results.some((existing) => existing.value === match.value)) {
+                results.push(match);
+            }
+        }
     }
     return {
         matches: results,
         classifier: classifierOutput
     };
 }
+export function privacyFilterOutputToMatches(text, output) {
+    if (!Array.isArray(output))
+        return [];
+    const matches = [];
+    for (const candidate of output) {
+        if (!candidate || typeof candidate !== 'object')
+            continue;
+        const span = candidate;
+        const rawLabel = span.entity_group ?? span.entity;
+        if (typeof rawLabel !== 'string')
+            continue;
+        const label = rawLabel.replace(/^[BIES]-/, '').toLowerCase();
+        const type = PRIVACY_FILTER_TYPE_MAP[label];
+        if (!type)
+            continue;
+        let value;
+        if (Number.isInteger(span.start) &&
+            Number.isInteger(span.end) &&
+            span.start >= 0 &&
+            span.end > span.start &&
+            span.end <= text.length) {
+            value = text.slice(span.start, span.end);
+        }
+        else if (typeof span.word === 'string') {
+            value = span.word.trim();
+        }
+        if (!value || !text.includes(value))
+            continue;
+        if (!matches.some((existing) => existing.value === value)) {
+            matches.push({ type, value });
+        }
+    }
+    return matches;
+}

package/dist/pii/storage.d.ts

@@ -27,7 +27,13 @@ export type RedisPiiVaultClient = {
 export type RedisPiiVaultStorageOptions = {
     keyPrefix?: string;
     tokenIndexKey?: string;
+    /** Expire Redis vault keys after this many seconds. Omit to keep them indefinitely. */
     ttlSeconds?: number;
+    /**
+     * Keep a process-local mirror and use it when a Redis operation fails.
+     * Disabled by default so Redis failures remain visible to callers.
+     */
+    fallbackToMemory?: boolean;
 };
 export declare function createRedisPiiVaultStorage(redis: RedisPiiVaultClient, options?: RedisPiiVaultStorageOptions): PiiVaultStorage;
 export declare class InMemoryPiiVaultStorage implements PiiVaultStorage {

package/dist/pii/storage.js

@@ -4,6 +4,13 @@ export function createPiiVaultStorage(adapter) {
 export function createRedisPiiVaultStorage(redis, options = {}) {
     const keyPrefix = options.keyPrefix ?? 'genkit-guard:pii';
     const tokenIndexKey = options.tokenIndexKey ?? `${keyPrefix}:tokens`;
+    const ttlSeconds = options.ttlSeconds;
+    if (ttlSeconds !== undefined && (!Number.isInteger(ttlSeconds) || ttlSeconds <= 0)) {
+        throw new Error('Redis PII vault ttlSeconds must be a positive integer.');
+    }
+    if (ttlSeconds !== undefined && !redis.expire) {
+        throw new Error('Redis PII vault ttlSeconds requires an expire method on the Redis client.');
+    }
     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);
@@ -11,31 +18,81 @@ export function createRedisPiiVaultStorage(redis, options = {}) {
         throw new Error('Redis PII vault storage requires hGet/hSet/hGetAll or hget/hset/hgetall methods.');
     }
     const scopeKey = (scopeId) => `${keyPrefix}:scope:${scopeId}`;
+    const fallback = options.fallbackToMemory ? new ExpiringInMemoryPiiVaultStorage(ttlSeconds) : undefined;
     async function maybeExpire(key) {
-        if (options.ttlSeconds && redis.expire) {
-            await redis.expire(key, options.ttlSeconds);
+        if (ttlSeconds !== undefined) {
+            await redis.expire(key, ttlSeconds);
+        }
+    }
+    async function withFallback(redisOperation, memoryOperation) {
+        try {
+            return await redisOperation();
+        }
+        catch (error) {
+            if (!fallback)
+                throw error;
+            return memoryOperation();
         }
     }
     return createPiiVaultStorage({
         async get(scopeId, token) {
-            return (await hGet(scopeKey(scopeId), token)) ?? undefined;
+            return withFallback(async () => (await hGet(scopeKey(scopeId), token)) ?? undefined, () => fallback.get(scopeId, token));
         },
         async getByToken(token) {
-            return (await hGet(tokenIndexKey, token)) ?? undefined;
+            return withFallback(async () => (await hGet(tokenIndexKey, token)) ?? undefined, () => fallback.getByToken(token));
         },
         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);
+            // Warm the opt-in fallback on every write so data written before an outage is available.
+            await fallback?.set(scopeId, token, value);
+            await withFallback(async () => {
+                await hSet(scopedKey, token, value);
+                await hSet(tokenIndexKey, token, value);
+                await maybeExpire(scopedKey);
+                await maybeExpire(tokenIndexKey);
+            }, () => undefined);
         },
         async entries(scopeId) {
-            const values = await hGetAll(scopeKey(scopeId));
-            return Object.entries(values).map(([token, value]) => ({ token, value }));
+            return withFallback(async () => {
+                const values = await hGetAll(scopeKey(scopeId));
+                return Object.entries(values).map(([token, value]) => ({ token, value }));
+            }, () => fallback.entries(scopeId));
         },
     });
 }
+class ExpiringInMemoryPiiVaultStorage {
+    ttlSeconds;
+    storage = new InMemoryPiiVaultStorage();
+    scopeExpiresAt = new Map();
+    tokenIndexExpiresAt;
+    constructor(ttlSeconds) {
+        this.ttlSeconds = ttlSeconds;
+    }
+    get(scopeId, token) {
+        return this.isScopeExpired(scopeId) ? undefined : this.storage.get(scopeId, token);
+    }
+    getByToken(token) {
+        return this.isTokenIndexExpired() ? undefined : this.storage.getByToken(token);
+    }
+    set(scopeId, token, value) {
+        this.storage.set(scopeId, token, value);
+        if (this.ttlSeconds !== undefined) {
+            const expiry = Date.now() + this.ttlSeconds * 1_000;
+            this.scopeExpiresAt.set(scopeId, expiry);
+            this.tokenIndexExpiresAt = expiry;
+        }
+    }
+    entries(scopeId) {
+        return this.isScopeExpired(scopeId) ? [] : this.storage.entries(scopeId);
+    }
+    isScopeExpired(scopeId) {
+        const expiry = this.scopeExpiresAt.get(scopeId);
+        return expiry !== undefined && expiry <= Date.now();
+    }
+    isTokenIndexExpired() {
+        return this.tokenIndexExpiresAt !== undefined && this.tokenIndexExpiresAt <= Date.now();
+    }
+}
 export class InMemoryPiiVaultStorage {
     scopes = new Map();
     tokenIndex = new Map();

package/dist/util/singleton.d.ts

@@ -1,7 +1,8 @@
+export declare const PRIVACY_FILTER_PIPELINE_TASK: "token-classification";
 export declare class ModelSingleton {
     private static extractors;
     private static nerClassifiers;
-    private static textClassifiers;
+    private static privacyFilters;
     static init(): void;
     static getExtractor(modelName?: string): Promise<any>;
     static getNER(modelName?: string): Promise<any>;

package/dist/util/singleton.js

@@ -3,12 +3,13 @@ import path from 'path';
 import fs from 'fs';
 import { fileURLToPath } from 'url';
 const __dirname = path.dirname(fileURLToPath(import.meta.url));
+export const PRIVACY_FILTER_PIPELINE_TASK = 'token-classification';
 env.allowRemoteModels = false;
 env.localModelPath = path.join(__dirname, '../../models');
 export class ModelSingleton {
     static extractors = new Map();
     static nerClassifiers = new Map();
-    static textClassifiers = new Map();
+    static privacyFilters = new Map();
     static init() {
         // Always resolve model path relative to the client app, not the library
         const projectRoot = process.cwd();
@@ -55,12 +56,14 @@ export class ModelSingleton {
         return this.nerClassifiers.get(modelName);
     }
     static async getPIIClassifier(modelName = 'openai/privacy-filter') {
-        if (!this.textClassifiers.has(modelName)) {
+        if (!this.privacyFilters.has(modelName)) {
             this.init();
-            const inst = await pipeline('text-classification', modelName);
-            this.textClassifiers.set(modelName, inst);
+            const inst = await pipeline(PRIVACY_FILTER_PIPELINE_TASK, modelName, {
+                dtype: 'q4',
+            });
+            this.privacyFilters.set(modelName, inst);
         }
-        return this.textClassifiers.get(modelName);
+        return this.privacyFilters.get(modelName);
     }
     static async preload(models) {
         const tasks = [];

package/package.json

@@ -22,7 +22,7 @@
     "huggingface"
   ],
   "license": "Apache-2.0",
-  "version": "0.0.11",
+  "version": "0.0.13",
   "type": "module",
   "exports": "./dist/index.js",
   "types": "./dist/index.d.ts",
@@ -34,9 +34,11 @@
   ],
   "scripts": {
     "prepare-models": "node scripts/download-model.js",
-    "build": "tsc",
-    "test": "npm run test:types && npm run test:storage && npm run test:concurrency",
+    "clean": "node -e \"require('node:fs').rmSync('dist', { recursive: true, force: true })\"",
+    "build": "npm run clean && tsc",
+    "test": "npm run test:types && npm run test:privacy-filter && 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:privacy-filter": "npm run build && node scripts/test-privacy-filter.js",
     "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",

package/scripts/test-privacy-filter.js

@@ -0,0 +1,116 @@
+import assert from 'node:assert/strict';
+import { detectPII, privacyFilterOutputToMatches } from '../dist/pii/detector.js';
+import { ModelSingleton, PRIVACY_FILTER_PIPELINE_TASK } from '../dist/util/singleton.js';
+import { InMemoryPiiVaultStorage } from '../dist/pii/storage.js';
+import { PiiTokenizer } from '../dist/pii/tokenizer.js';
+import { guard } from '../dist/index.js';
+
+assert.equal(PRIVACY_FILTER_PIPELINE_TASK, 'token-classification');
+
+const input = 'Contact Alice Smith at alice@example.com. Her key is sk-live-secret.';
+const expectedSpans = [
+  { entity_group: 'private_person', word: ' Alice Smith' },
+  { entity_group: 'private_email', word: ' alice@example.com' },
+  { entity_group: 'secret', word: ' sk-live-secret' },
+];
+
+assert.deepEqual(privacyFilterOutputToMatches(input, expectedSpans), [
+  { type: 'NAME', value: 'Alice Smith' },
+  { type: 'EMAIL', value: 'alice@example.com' },
+  { type: 'SECRET', value: 'sk-live-secret' },
+]);
+
+assert.deepEqual(
+  privacyFilterOutputToMatches('Home: 10 Green Street', [
+    { entity_group: 'private_address', start: 6, end: 21, word: 'wrong fallback' },
+    { entity_group: 'O', word: 'Home' },
+    { entity_group: 'private_date', word: 'not present' },
+    null,
+  ]),
+  [{ type: 'ADDRESS', value: '10 Green Street' }]
+);
+assert.deepEqual(privacyFilterOutputToMatches(input, { invalid: true }), []);
+
+const originalGetPIIClassifier = ModelSingleton.getPIIClassifier;
+const originalGetExtractor = ModelSingleton.getExtractor;
+let receivedOptions;
+ModelSingleton.getPIIClassifier = async () => async (_text, options) => {
+  receivedOptions = options;
+  return expectedSpans;
+};
+
+try {
+  const detection = await detectPII(input, { mode: 'classifier' });
+  assert.deepEqual(receivedOptions, { aggregation_strategy: 'simple' });
+
+  // The regex email and model email are deduplicated; model-only name and secret become matches.
+  assert.deepEqual(detection.matches, [
+    { type: 'EMAIL', value: 'alice@example.com' },
+    { type: 'NAME', value: 'Alice Smith' },
+    { type: 'SECRET', value: 'sk-live-secret' },
+  ]);
+
+  const storage = new InMemoryPiiVaultStorage();
+  const tokenizer = new PiiTokenizer({ scopeId: 'privacy-filter-test', storage });
+  const masked = await tokenizer.mask(input, detection.matches);
+
+  assert.doesNotMatch(masked.maskedText, /Alice Smith|alice@example\.com|sk-live-secret/);
+  assert.match(masked.maskedText, /\[\[NAME_/);
+  assert.match(masked.maskedText, /\[\[EMAIL_/);
+  assert.match(masked.maskedText, /\[\[SECRET_/);
+  assert.equal(await tokenizer.unmask(masked.maskedText), input);
+} finally {
+  ModelSingleton.getPIIClassifier = originalGetPIIClassifier;
+}
+
+// Reproduce a Genkit multi-turn response: an inner turn creates the token in one scope, while
+// the final response passes through middleware attached to another context and scope.
+const crossTurnStorage = new InMemoryPiiVaultStorage();
+const innerTurn = new PiiTokenizer({ scopeId: 'inner-turn', storage: crossTurnStorage });
+const crossTurnMasked = await innerTurn.mask('owner@example.com', [
+  { type: 'EMAIL', value: 'owner@example.com' },
+]);
+
+ModelSingleton.getExtractor = async () => async () => ({ tolist: () => [] });
+ModelSingleton.getPIIClassifier = async () => async () => [];
+
+try {
+  const middleware = guard({
+    intent: { semantic: { threshold: 0, intents: {} } },
+    pii: {
+      mode: 'classifier',
+      vault: { storage: crossTurnStorage, scopeId: 'outer-turn' },
+    },
+    logging: { enabled: false },
+  });
+
+  const response = await middleware.model(
+    { prompt: 'Fetch blob metadata' },
+    {},
+    async () => ({ answer: `File owner: ${crossTurnMasked.maskedText}` })
+  );
+  assert.deepEqual(response, { answer: 'File owner: owner@example.com' });
+
+  const unknownTokenResponse = await middleware.model(
+    { prompt: 'Fetch blob metadata' },
+    {},
+    async () => ({ answer: '[[EMAIL_unknownnamespace_99]]' })
+  );
+  assert.deepEqual(unknownTokenResponse, { answer: '[[EMAIL_unknownnamespace_99]]' });
+
+  let toolInput;
+  await middleware.tool(
+    { toolRequest: { name: 'sendEmail', input: { recipient: crossTurnMasked.maskedText } } },
+    {},
+    async (request) => {
+      toolInput = request.toolRequest.input;
+      return { sent: true };
+    }
+  );
+  assert.deepEqual(toolInput, { recipient: 'owner@example.com' });
+} finally {
+  ModelSingleton.getExtractor = originalGetExtractor;
+  ModelSingleton.getPIIClassifier = originalGetPIIClassifier;
+}
+
+console.log('Privacy Filter pipeline, masking and cross-turn unmasking tests passed.');

package/scripts/test-storage.js

@@ -1,78 +1,137 @@
-import assert from 'node:assert/strict';
-import {
-  createPiiVaultStorage,
-  createRedisPiiVaultStorage,
-  InMemoryPiiVaultStorage,
-} from '../dist/index.js';
-import { PiiTokenizer } from '../dist/pii/tokenizer.js';
-
-class FakeRedis {
-  hashes = new Map();
-  expirations = new Map();
-
-  async hGet(key, field) {
-    return this.hashes.get(key)?.[field] ?? null;
-  }
-
-  async hSet(key, field, value) {
-    const hash = this.hashes.get(key) ?? {};
-    hash[field] = value;
-    this.hashes.set(key, hash);
-  }
-
-  async hGetAll(key) {
-    return this.hashes.get(key) ?? {};
-  }
-
-  async expire(key, seconds) {
-    this.expirations.set(key, seconds);
-  }
-}
-
-const memory = new InMemoryPiiVaultStorage();
-await memory.set('scope-a', '[[EMAIL_token_0]]', 'a@example.com');
-assert.equal(await memory.get('scope-a', '[[EMAIL_token_0]]'), 'a@example.com');
-assert.equal(await memory.get('scope-b', '[[EMAIL_token_0]]'), undefined);
-assert.equal(await memory.getByToken('[[EMAIL_token_0]]'), 'a@example.com');
-
-const custom = createPiiVaultStorage({
-  async get(scopeId, token) {
-    return scopeId === 'custom' && token === '[[EMAIL_custom_0]]' ? 'custom@example.com' : undefined;
-  },
-  async set() {},
-  async entries(scopeId) {
-    return scopeId === 'custom'
-      ? [{ token: '[[EMAIL_custom_0]]', value: 'custom@example.com' }]
-      : [];
-  },
-});
-
-const customTokenizer = new PiiTokenizer({ scopeId: 'custom', storage: custom });
-assert.equal(await customTokenizer.unmask('Hi [[EMAIL_custom_0]]'), 'Hi custom@example.com');
-
-const redis = new FakeRedis();
-const redisStorage = createRedisPiiVaultStorage(redis, {
-  keyPrefix: 'test:pii',
-  ttlSeconds: 60,
-});
-
-const userATokenizer = new PiiTokenizer({ scopeId: 'tenant:userA', storage: redisStorage });
-const userBTokenizer = new PiiTokenizer({ scopeId: 'tenant:userB', storage: redisStorage });
-
-const userAMasked = await userATokenizer.mask('Email alice@example.com', [
-  { type: 'EMAIL', value: 'alice@example.com' },
-]);
-const userBMasked = await userBTokenizer.mask('Email bob@example.com', [
-  { type: 'EMAIL', value: 'bob@example.com' },
-]);
-
-assert.notEqual(userAMasked.maskedText, userBMasked.maskedText);
-assert.equal(await userATokenizer.unmask(userAMasked.maskedText), 'Email alice@example.com');
-assert.equal(await userBTokenizer.unmask(userBMasked.maskedText), 'Email bob@example.com');
-assert.equal(await userATokenizer.unmask(userBMasked.maskedText), userBMasked.maskedText);
-
-const laterPassTokenizer = new PiiTokenizer({ scopeId: 'tenant:userA:later', storage: redisStorage });
-await laterPassTokenizer.importTokens(userAMasked.maskedText);
-assert.equal(await laterPassTokenizer.unmask(userAMasked.maskedText), 'Email alice@example.com');
-
-console.log('PII vault storage tests passed.');
+import assert from 'node:assert/strict';
+import {
+  createPiiVaultStorage,
+  createRedisPiiVaultStorage,
+  InMemoryPiiVaultStorage,
+} from '../dist/index.js';
+import { PiiTokenizer } from '../dist/pii/tokenizer.js';
+
+class FakeRedis {
+  hashes = new Map();
+  expirations = new Map();
+
+  async hGet(key, field) {
+    return this.hashes.get(key)?.[field] ?? null;
+  }
+
+  async hSet(key, field, value) {
+    const hash = this.hashes.get(key) ?? {};
+    hash[field] = value;
+    this.hashes.set(key, hash);
+  }
+
+  async hGetAll(key) {
+    return this.hashes.get(key) ?? {};
+  }
+
+  async expire(key, seconds) {
+    this.expirations.set(key, seconds);
+  }
+}
+
+const memory = new InMemoryPiiVaultStorage();
+await memory.set('scope-a', '[[EMAIL_token_0]]', 'a@example.com');
+assert.equal(await memory.get('scope-a', '[[EMAIL_token_0]]'), 'a@example.com');
+assert.equal(await memory.get('scope-b', '[[EMAIL_token_0]]'), undefined);
+assert.equal(await memory.getByToken('[[EMAIL_token_0]]'), 'a@example.com');
+
+const custom = createPiiVaultStorage({
+  async get(scopeId, token) {
+    return scopeId === 'custom' && token === '[[EMAIL_custom_0]]' ? 'custom@example.com' : undefined;
+  },
+  async set() {},
+  async entries(scopeId) {
+    return scopeId === 'custom'
+      ? [{ token: '[[EMAIL_custom_0]]', value: 'custom@example.com' }]
+      : [];
+  },
+});
+
+const customTokenizer = new PiiTokenizer({ scopeId: 'custom', storage: custom });
+assert.equal(await customTokenizer.unmask('Hi [[EMAIL_custom_0]]'), 'Hi custom@example.com');
+
+const redis = new FakeRedis();
+const redisStorage = createRedisPiiVaultStorage(redis, {
+  keyPrefix: 'test:pii',
+  ttlSeconds: 60,
+});
+
+const userATokenizer = new PiiTokenizer({ scopeId: 'tenant:userA', storage: redisStorage });
+const userBTokenizer = new PiiTokenizer({ scopeId: 'tenant:userB', storage: redisStorage });
+
+const userAMasked = await userATokenizer.mask('Email alice@example.com', [
+  { type: 'EMAIL', value: 'alice@example.com' },
+]);
+const userBMasked = await userBTokenizer.mask('Email bob@example.com', [
+  { type: 'EMAIL', value: 'bob@example.com' },
+]);
+
+assert.notEqual(userAMasked.maskedText, userBMasked.maskedText);
+assert.equal(await userATokenizer.unmask(userAMasked.maskedText), 'Email alice@example.com');
+assert.equal(await userBTokenizer.unmask(userBMasked.maskedText), 'Email bob@example.com');
+assert.equal(await userATokenizer.unmask(userBMasked.maskedText), userBMasked.maskedText);
+
+const laterPassTokenizer = new PiiTokenizer({ scopeId: 'tenant:userA:later', storage: redisStorage });
+await laterPassTokenizer.importTokens(userAMasked.maskedText);
+assert.equal(await laterPassTokenizer.unmask(userAMasked.maskedText), 'Email alice@example.com');
+
+assert.equal(redis.expirations.get('test:pii:scope:tenant:userA'), 60);
+assert.equal(redis.expirations.get('test:pii:tokens'), 60);
+
+class FailingRedis extends FakeRedis {
+  failed = false;
+
+  fail() {
+    this.failed = true;
+  }
+
+  async hGet(key, field) {
+    if (this.failed) throw new Error('Redis unavailable');
+    return super.hGet(key, field);
+  }
+
+  async hSet(key, field, value) {
+    if (this.failed) throw new Error('Redis unavailable');
+    return super.hSet(key, field, value);
+  }
+
+  async hGetAll(key) {
+    if (this.failed) throw new Error('Redis unavailable');
+    return super.hGetAll(key);
+  }
+}
+
+const failingRedis = new FailingRedis();
+const resilientStorage = createRedisPiiVaultStorage(failingRedis, { fallbackToMemory: true });
+await resilientStorage.set('resilient-scope', '[[EMAIL_resilient_0]]', 'safe@example.com');
+failingRedis.fail();
+assert.equal(
+  await resilientStorage.get('resilient-scope', '[[EMAIL_resilient_0]]'),
+  'safe@example.com'
+);
+assert.equal(await resilientStorage.getByToken('[[EMAIL_resilient_0]]'), 'safe@example.com');
+assert.deepEqual(await resilientStorage.entries('resilient-scope'), [
+  { token: '[[EMAIL_resilient_0]]', value: 'safe@example.com' },
+]);
+await resilientStorage.set('resilient-scope', '[[PHONE_resilient_1]]', '0400000000');
+assert.equal(await resilientStorage.getByToken('[[PHONE_resilient_1]]'), '0400000000');
+
+const strictRedis = new FailingRedis();
+const strictStorage = createRedisPiiVaultStorage(strictRedis);
+strictRedis.fail();
+await assert.rejects(() => strictStorage.get('scope', 'token'), /Redis unavailable/);
+
+assert.throws(
+  () => createRedisPiiVaultStorage(new FakeRedis(), { ttlSeconds: 0 }),
+  /positive integer/
+);
+assert.throws(
+  () =>
+    createRedisPiiVaultStorage(
+      { hGet: async () => null, hSet: async () => undefined, hGetAll: async () => ({}) },
+      { ttlSeconds: 60 }
+    ),
+  /requires an expire method/
+);
+
+console.log('PII vault storage tests passed.');

package/scripts/test-types.ts

@@ -1,21 +1,21 @@
-import { guard, createRedisPiiVaultStorage, type RedisPiiVaultClient } from '../src/index.js';
-
-const redis: RedisPiiVaultClient = {
-  async hGet() {
-    return undefined;
-  },
-  async hSet() {},
-  async hGetAll() {
-    return {};
-  },
-};
-
-guard({
-  pii: {
-    reversible: true,
-    vault: {
-      storage: createRedisPiiVaultStorage(redis),
-      scopeId: (req: any, ctx: any) => ctx?.auth?.sessionId ?? req?.metadata?.requestId,
-    },
-  },
-});
+import { guard, createRedisPiiVaultStorage, type RedisPiiVaultClient } from '../src/index.js';
+
+const redis: RedisPiiVaultClient = {
+  async hGet() {
+    return undefined;
+  },
+  async hSet() {},
+  async hGetAll() {
+    return {};
+  },
+};
+
+guard({
+  pii: {
+    reversible: true,
+    vault: {
+      storage: createRedisPiiVaultStorage(redis),
+      scopeId: (req: any, ctx: any) => ctx?.auth?.sessionId ?? req?.metadata?.requestId,
+    },
+  },
+});