@privacyscrubber/sdk 2.3.7 → 2.4.0

package/fastmcp.cjs

@@ -26,6 +26,204 @@ const ZTDS_PATTERNS = {
 
 const TOKEN_PATTERN = /\[([A-Z_]+)_TOKEN_([a-zA-Z0-9_-]+)\]/g;
 
+function escapeRegex(s) {
+  return s.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
+}
+
+function buildStreamReplacer(tokenMap = {}, options = {}) {
+  if (!tokenMap || typeof tokenMap !== 'object' || Object.keys(tokenMap).length === 0) {
+    return (text) => text || '';
+  }
+
+  const lookup = new Map();
+  const patterns = [];
+  const sortedKeys = Object.keys(tokenMap).sort((a, b) => b.length - a.length);
+
+  for (const k of sortedKeys) {
+    const val = tokenMap[k];
+    const inner = k.replace(/^\[|<|\{\{|__|\]|>|\}\}|__/g, '');
+    const match = inner.match(/^([A-Za-z_0-9]+?)[-_ ]?(\d+)$/);
+
+    if (match) {
+      const label = escapeRegex(match[1]);
+      const idx = match[2];
+      const patternStr = '(?:\\[\\s*' + label + '[-_\\s]*0*' + idx + '\\s*\\]|<\\s*' + label + '[-_\\s]*0*' + idx + '\\s*>|\\{\\{\\s*' + label + '[-_\\s]*0*' + idx + '\\s*\\}\\}|__\\s*' + label + '[-_\\s]*0*' + idx + '\\s*__|\\b' + label + '[-_\\s]+0*' + idx + '\\b)';
+      patterns.push(patternStr);
+      const normKey = (match[1] + '_' + parseInt(idx, 10)).toUpperCase();
+      lookup.set(normKey, val);
+      lookup.set(k, val);
+      lookup.set(k.toUpperCase(), val);
+    } else {
+      const safe = escapeRegex(k);
+      patterns.push(safe);
+      lookup.set(k, val);
+      lookup.set(k.toUpperCase(), val);
+    }
+  }
+
+  const unifiedRegex = new RegExp('(?:' + patterns.join('|') + ')(\'s|’s|s)?', 'gi');
+
+  return function replaceTokens(text) {
+    if (!text) return '';
+    unifiedRegex.lastIndex = 0;
+    return text.replace(unifiedRegex, (match, suffix) => {
+      const isPossessive = suffix && /('s|’s)$/i.test(suffix);
+      const coreMatch = suffix ? match.slice(0, match.length - suffix.length) : match;
+      let restored = null;
+
+      if (lookup.has(coreMatch)) {
+        restored = lookup.get(coreMatch);
+      } else if (lookup.has(coreMatch.toUpperCase())) {
+        restored = lookup.get(coreMatch.toUpperCase());
+      } else {
+        const inner = coreMatch.replace(/^\[|<|\{\{|__|\]|>|\}\}|__/g, '').trim();
+        const m = inner.match(/^([A-Za-z_0-9]+?)[-_ ]?(\d+)$/);
+        if (m) {
+          const normKey = (m[1] + '_' + parseInt(m[2], 10)).toUpperCase();
+          if (lookup.has(normKey)) {
+            restored = lookup.get(normKey);
+          }
+        }
+      }
+
+      if (restored !== null) {
+        return restored + (isPossessive ? "'s" : '');
+      }
+
+      if (lookup.has(match)) {
+        return lookup.get(match);
+      }
+
+      return match;
+    });
+  };
+}
+
+class StreamingDetokenizer {
+  constructor(tokenMap = {}, options = {}) {
+    this.tokenMap = tokenMap && typeof tokenMap === 'object' ? tokenMap : {};
+    this.options = options || {};
+    this.replacer = buildStreamReplacer(this.tokenMap, this.options);
+    this.pendingBuffer = '';
+    this.hasTokens = Object.keys(this.tokenMap).length > 0;
+    this._lastTokenCount = Object.keys(this.tokenMap).length;
+    this.decoder = new TextDecoder('utf-8', { fatal: false });
+  }
+
+  get buffer() {
+    return this.pendingBuffer;
+  }
+
+  set buffer(val) {
+    this.pendingBuffer = typeof val === 'string' ? val : '';
+  }
+
+  updateTokenMap(newMap) {
+    if (newMap && typeof newMap === 'object') {
+      this.tokenMap = newMap;
+      this.replacer = buildStreamReplacer(this.tokenMap, this.options);
+      this.hasTokens = Object.keys(this.tokenMap).length > 0;
+      this._lastTokenCount = Object.keys(this.tokenMap).length;
+    }
+  }
+
+  reset() {
+    this.pendingBuffer = '';
+  }
+
+  feed(chunk) {
+    if (chunk === undefined || chunk === null) return '';
+    let text = '';
+    if (typeof chunk === 'string') {
+      text = chunk;
+    } else if (chunk instanceof Uint8Array || (typeof Buffer !== 'undefined' && Buffer.isBuffer(chunk))) {
+      text = this.decoder.decode(chunk, { stream: true });
+    } else {
+      text = String(chunk);
+    }
+
+    if (!text) return '';
+
+    const currentTokenCount = Object.keys(this.tokenMap).length;
+    if (currentTokenCount !== this._lastTokenCount) {
+      this.replacer = buildStreamReplacer(this.tokenMap, this.options);
+      this.hasTokens = currentTokenCount > 0;
+      this._lastTokenCount = currentTokenCount;
+    }
+
+    if (!this.hasTokens) return text;
+
+    this.pendingBuffer += text;
+
+    const tailLen = 160;
+    const tail = this.pendingBuffer.length > tailLen ? this.pendingBuffer.slice(-tailLen) : this.pendingBuffer;
+    const offset = this.pendingBuffer.length > tailLen ? this.pendingBuffer.length - tailLen : 0;
+    const m = tail.match(/(?:\[|<|\{\{?|__?)[A-Za-z0-9_\-А-Яа-яЁё ]{0,128}$/);
+
+    if (m) {
+      const candidate = m[0];
+      const isClosedUnderscoreToken = candidate.length >= 4 && candidate.startsWith('__') && candidate.endsWith('__');
+      if (isClosedUnderscoreToken && this.replacer(candidate) !== candidate) {
+        const toEmit = this.pendingBuffer;
+        this.pendingBuffer = '';
+        return this.replacer(toEmit);
+      }
+      const splitIdx = offset + m.index;
+      const safeToEmit = this.pendingBuffer.slice(0, splitIdx);
+      this.pendingBuffer = this.pendingBuffer.slice(splitIdx);
+      return safeToEmit ? this.replacer(safeToEmit) : '';
+    }
+
+    const toEmit = this.pendingBuffer;
+    this.pendingBuffer = '';
+    return this.replacer(toEmit);
+  }
+
+  process(chunk) {
+    return this.feed(chunk);
+  }
+
+  flush() {
+    let text = '';
+    try {
+      text = this.decoder.decode(new Uint8Array(), { stream: false });
+    } catch (_) {}
+
+    if (text) {
+      this.pendingBuffer += text;
+    }
+
+    if (!this.pendingBuffer) return '';
+    const remaining = this.pendingBuffer;
+    this.pendingBuffer = '';
+    return this.replacer(remaining);
+  }
+
+  async* wrapStream(asyncIterable) {
+    for await (const chunk of asyncIterable) {
+      if (chunk === null || chunk === undefined) continue;
+      let text = '';
+      if (typeof chunk === 'string') {
+        text = chunk;
+      } else if (typeof chunk?.text === 'string') {
+        text = chunk.text;
+      } else if (typeof chunk?.content === 'string') {
+        text = chunk.content;
+      } else if (chunk instanceof Uint8Array || (typeof Buffer !== 'undefined' && Buffer.isBuffer(chunk))) {
+        text = chunk;
+      } else {
+        text = String(chunk);
+      }
+
+      const emitted = this.feed(text);
+      if (emitted) yield emitted;
+    }
+
+    const finalChunk = this.flush();
+    if (finalChunk) yield finalChunk;
+  }
+}
+
 class ZTDSFastMCPGuard {
   constructor(options = {}) {
     this.enabledEntities = options.enabledEntities || Object.keys(ZTDS_PATTERNS);
@@ -33,8 +231,10 @@ class ZTDSFastMCPGuard {
     this.sanitizeOutputs = options.sanitizeOutputs !== false;
     this.useCryptoTokens = options.useCryptoTokens !== false;
     this.autoZeroize = options.autoZeroize !== false;
+    this.maxSessionTokens = options.maxSessionTokens || 5000;
     this._sessionMaps = new Map();
     this._entityMaps = new Map();
+    this._streamDetokenizers = new Map();
   }
 
   _generateToken(entityType, tokenMap) {
@@ -86,6 +286,14 @@ class ZTDSFastMCPGuard {
         if (entityMap.has(match)) {
           return entityMap.get(match);
         }
+        if (tokenMap.size >= this.maxSessionTokens) {
+          const oldestToken = tokenMap.keys().next().value;
+          const oldestMatch = tokenMap.get(oldestToken);
+          tokenMap.delete(oldestToken);
+          if (oldestMatch && entityMap.has(oldestMatch)) {
+            entityMap.delete(oldestMatch);
+          }
+        }
         const token = this._generateToken(entityType, tokenMap);
         tokenMap.set(token, match);
         entityMap.set(match, token);
@@ -148,6 +356,65 @@ class ZTDSFastMCPGuard {
     return val;
   }
 
+  getStreamDetokenizer(sessionId) {
+    if (!this._streamDetokenizers.has(sessionId)) {
+      const mapProxy = this._sessionMaps.get(sessionId) || new Map();
+      const objMap = {};
+      for (const [k, v] of mapProxy.entries()) {
+        objMap[k] = v;
+      }
+      this._streamDetokenizers.set(sessionId, new StreamingDetokenizer(objMap));
+    } else {
+      const mapProxy = this._sessionMaps.get(sessionId);
+      if (mapProxy) {
+        const objMap = {};
+        for (const [k, v] of mapProxy.entries()) {
+          objMap[k] = v;
+        }
+        this._streamDetokenizers.get(sessionId).updateTokenMap(objMap);
+      }
+    }
+    return this._streamDetokenizers.get(sessionId);
+  }
+
+  feedStream(chunk, sessionId) {
+    return this.getStreamDetokenizer(sessionId).feed(chunk);
+  }
+
+  flushStream(sessionId) {
+    const detok = this._streamDetokenizers.get(sessionId);
+    return detok ? detok.flush() : "";
+  }
+
+  async *restoreStream(asyncIterable, sessionId) {
+    const detok = this.getStreamDetokenizer(sessionId);
+    try {
+      for await (const chunk of asyncIterable) {
+        if (chunk === null || chunk === undefined) continue;
+        let text = "";
+        if (typeof chunk === "string") {
+          text = chunk;
+        } else if (typeof chunk?.text === "string") {
+          text = chunk.text;
+        } else if (typeof chunk?.content === "string") {
+          text = chunk.content;
+        } else if (chunk instanceof Uint8Array || (typeof Buffer !== "undefined" && Buffer.isBuffer(chunk))) {
+          text = chunk;
+        } else {
+          text = String(chunk);
+        }
+        const emitted = detok.feed(text);
+        if (emitted) yield emitted;
+      }
+      const finalChunk = detok.flush();
+      if (finalChunk) yield finalChunk;
+    } finally {
+      if (this.autoZeroize) {
+        this.zeroizeSession(sessionId);
+      }
+    }
+  }
+
   wrapTool(toolName, handler) {
     return async (args, context = {}) => {
       const sessionId =
@@ -160,6 +427,7 @@ class ZTDSFastMCPGuard {
         unmask: (val) => this.restoreObject(val, sessionId),
       };
 
+      let isStreaming = false;
       try {
         let processedArgs = args;
         if (this.sanitizeInputs && args !== undefined && args !== null) {
@@ -168,6 +436,11 @@ class ZTDSFastMCPGuard {
 
         const rawResult = await handler(processedArgs, enrichedContext);
 
+        if (rawResult && typeof rawResult[Symbol.asyncIterator] === "function") {
+          isStreaming = true;
+          return this.restoreStream(rawResult, sessionId);
+        }
+
         let finalResult = rawResult;
         if (this.sanitizeOutputs && rawResult !== undefined && rawResult !== null) {
           if (typeof rawResult === "string") {
@@ -176,12 +449,22 @@ class ZTDSFastMCPGuard {
             finalResult = {
               ...rawResult,
               content: rawResult.content.map((block) => {
-                if (block && block.type === "text" && typeof block.text === "string") {
+                if (!block || typeof block !== "object") return block;
+                if (block.type === "text" && typeof block.text === "string") {
                   return {
                     ...block,
                     text: this.sanitizeText(block.text, sessionId).sanitized,
                   };
                 }
+                if ((block.type === "resource" || block.type === "embeddedResource") && block.resource && typeof block.resource.text === "string") {
+                  return {
+                    ...block,
+                    resource: {
+                      ...block.resource,
+                      text: this.sanitizeText(block.resource.text, sessionId).sanitized,
+                    },
+                  };
+                }
                 return block;
               }),
             };
@@ -215,7 +498,7 @@ class ZTDSFastMCPGuard {
           _ztds: ztdsMeta,
         };
       } finally {
-        if (this.autoZeroize) {
+        if (!isStreaming && this.autoZeroize) {
           this.zeroizeSession(sessionId);
         }
       }
@@ -231,6 +514,10 @@ class ZTDSFastMCPGuard {
       this._entityMaps.get(sessionId).clear();
       this._entityMaps.delete(sessionId);
     }
+    if (this._streamDetokenizers.has(sessionId)) {
+      this._streamDetokenizers.get(sessionId).reset();
+      this._streamDetokenizers.delete(sessionId);
+    }
   }
 }
 
@@ -266,6 +553,8 @@ module.exports = {
   ZTDSFastMCPGuard,
   ZTDS_PATTERNS,
   TOKEN_PATTERN,
+  StreamingDetokenizer,
+  createStreamTokenRestorer: (tokenMap, options) => new StreamingDetokenizer(tokenMap, options),
   createZTDSGuard,
   ztdsGuardTool,
   wrapExecute,

package/fastmcp.d.ts

@@ -26,6 +26,19 @@ export interface ZTDSContext {
 export declare const ZTDS_PATTERNS: Record<string, RegExp>;
 export declare const TOKEN_PATTERN: RegExp;
 
+export declare class StreamingDetokenizer {
+  constructor(tokenMap?: Record<string, string>, options?: object);
+  buffer: string;
+  feed(chunk: string | Uint8Array | Buffer): string;
+  process(chunk: string | Uint8Array | Buffer): string;
+  flush(): string;
+  reset(): void;
+  updateTokenMap(newMap: Record<string, string>): void;
+  wrapStream<T = any>(asyncIterable: AsyncIterable<T>): AsyncGenerator<T>;
+}
+
+export declare function createStreamTokenRestorer(tokenMap: Record<string, string>, options?: object): StreamingDetokenizer;
+
 export declare class ZTDSFastMCPGuard {
   constructor(options?: ZTDSGuardOptions);
   enabledEntities: string[];
@@ -39,6 +52,11 @@ export declare class ZTDSFastMCPGuard {
   sanitizeObject(val: unknown, sessionId: string): unknown;
   restoreObject(val: unknown, sessionId: string): unknown;
 
+  getStreamDetokenizer(sessionId: string): StreamingDetokenizer;
+  feedStream(chunk: string, sessionId: string): string;
+  flushStream(sessionId: string): string;
+  restoreStream<T = any>(asyncIterable: AsyncIterable<T>, sessionId: string): AsyncGenerator<T>;
+
   wrapTool<TArgs = unknown, TResult = unknown>(
     toolName: string,
     handler: (args: TArgs, context: ZTDSContext) => Promise<TResult>

package/fastmcp.js

@@ -11,6 +11,8 @@ const fastmcpGuard = require('./fastmcp.cjs');
 export const ZTDSFastMCPGuard = fastmcpGuard.ZTDSFastMCPGuard;
 export const ZTDS_PATTERNS = fastmcpGuard.ZTDS_PATTERNS;
 export const TOKEN_PATTERN = fastmcpGuard.TOKEN_PATTERN;
+export const StreamingDetokenizer = fastmcpGuard.StreamingDetokenizer;
+export const createStreamTokenRestorer = fastmcpGuard.createStreamTokenRestorer;
 export const createZTDSGuard = fastmcpGuard.createZTDSGuard;
 export const ztdsGuardTool = fastmcpGuard.ztdsGuardTool;
 export const wrapExecute = fastmcpGuard.wrapExecute;
@@ -20,6 +22,8 @@ export default {
   ZTDSFastMCPGuard,
   ZTDS_PATTERNS,
   TOKEN_PATTERN,
+  StreamingDetokenizer,
+  createStreamTokenRestorer,
   createZTDSGuard,
   ztdsGuardTool,
   wrapExecute,