mcp-rubber-duck 1.7.0 → 1.9.0

package/src/guardrails/plugins/pii-redactor/index.ts

@@ -0,0 +1,203 @@
+import { BaseGuardrailPlugin } from '../base-plugin.js';
+import { GuardrailPhase, GuardrailContext, GuardrailResult } from '../../types.js';
+import { PIIRedactorConfig } from '../../../config/types.js';
+import { PIIDetector, PIIDetectorConfig } from './detectors.js';
+import { Pseudonymizer } from './pseudonymizer.js';
+import { logger } from '../../../utils/logger.js';
+
+/**
+ * PII Redactor plugin - detects and redacts sensitive data
+ */
+export class PIIRedactorPlugin extends BaseGuardrailPlugin {
+  name = 'pii_redactor';
+  phases: GuardrailPhase[] = ['pre_request', 'post_response', 'pre_tool_input', 'post_tool_output'];
+
+  private detector!: PIIDetector;
+  private pseudonymizer!: Pseudonymizer;
+  private restoreOnResponse: boolean = false;
+  private logDetections: boolean = true;
+
+  async initialize(config: Record<string, unknown>): Promise<void> {
+    await super.initialize(config);
+
+    const typedConfig = config as Partial<PIIRedactorConfig>;
+
+    const detectorConfig: PIIDetectorConfig = {
+      detectEmails: typedConfig.detect_emails ?? true,
+      detectPhones: typedConfig.detect_phones ?? true,
+      detectSSN: typedConfig.detect_ssn ?? true,
+      detectAPIKeys: typedConfig.detect_api_keys ?? true,
+      detectCreditCards: typedConfig.detect_credit_cards ?? true,
+      detectIPAddresses: typedConfig.detect_ip_addresses ?? false,
+      customPatterns: typedConfig.custom_patterns ?? [],
+      allowlist: typedConfig.allowlist ?? [],
+      allowlistDomains: typedConfig.allowlist_domains ?? [],
+    };
+
+    this.detector = new PIIDetector(detectorConfig);
+    this.pseudonymizer = new Pseudonymizer();
+    this.restoreOnResponse = typedConfig.restore_on_response ?? false;
+    this.logDetections = typedConfig.log_detections ?? true;
+    this.priority = typedConfig.priority ?? 25;
+  }
+
+  async execute(phase: GuardrailPhase, context: GuardrailContext): Promise<GuardrailResult> {
+    switch (phase) {
+      case 'pre_request':
+      case 'pre_tool_input':
+        return this.redactPII(context, phase);
+
+      case 'post_response':
+      case 'post_tool_output':
+        if (this.restoreOnResponse) {
+          return this.restorePII(context, phase);
+        }
+        return this.allow(context);
+
+      default:
+        return this.allow(context);
+    }
+  }
+
+  private redactPII(
+    context: GuardrailContext,
+    phase: GuardrailPhase
+  ): Promise<GuardrailResult> {
+    let textToScan: string;
+    let field: string;
+
+    if (phase === 'pre_request') {
+      textToScan = context.prompt || '';
+      field = 'prompt';
+    } else {
+      textToScan = JSON.stringify(context.toolArgs || {});
+      field = 'toolArgs';
+    }
+
+    if (!textToScan) {
+      return Promise.resolve(this.allow(context));
+    }
+
+    const detections = this.detector.detect(textToScan);
+
+    if (detections.length === 0) {
+      return Promise.resolve(this.allow(context));
+    }
+
+    // Log detections
+    if (this.logDetections) {
+      logger.info(`PII detected in ${field}:`, {
+        requestId: context.requestId,
+        types: [...new Set(detections.map((d) => d.type))],
+        count: detections.length,
+      });
+    }
+
+    // Pseudonymize
+    const { text: redactedText, mappings } = this.pseudonymizer.pseudonymize(
+      textToScan,
+      detections
+    );
+
+    // Store mappings for potential restoration
+    context.metadata.set('pii_mappings', mappings);
+
+    // Record modification
+    this.addModification(
+      context,
+      phase,
+      field,
+      `Redacted ${detections.length} PII items: ${[...new Set(detections.map((d) => d.type))].join(', ')}`,
+      undefined, // Don't store original (contains PII)
+      undefined  // Don't store new (would expose placeholder patterns)
+    );
+
+    // Apply modification
+    if (phase === 'pre_request') {
+      context.prompt = redactedText;
+      // Also update the last message if present
+      if (context.messages.length > 0) {
+        const lastIndex = context.messages.length - 1;
+        context.messages[lastIndex] = {
+          ...context.messages[lastIndex],
+          content: redactedText,
+        };
+      }
+    } else {
+      try {
+        context.toolArgs = JSON.parse(redactedText) as Record<string, unknown>;
+      } catch {
+        // If parse fails, store as string
+        context.toolArgs = { _redacted: redactedText };
+      }
+    }
+
+    return Promise.resolve(this.modify(context));
+  }
+
+  private restorePII(
+    context: GuardrailContext,
+    phase: GuardrailPhase
+  ): Promise<GuardrailResult> {
+    const mappings = context.metadata.get('pii_mappings') as Map<string, string> | undefined;
+
+    if (!mappings || mappings.size === 0) {
+      return Promise.resolve(this.allow(context));
+    }
+
+    let textToRestore: string;
+
+    if (phase === 'post_response') {
+      textToRestore = context.response || '';
+    } else {
+      textToRestore =
+        typeof context.toolResult === 'string'
+          ? context.toolResult
+          : JSON.stringify(context.toolResult);
+    }
+
+    if (!textToRestore) {
+      return Promise.resolve(this.allow(context));
+    }
+
+    const restoredText = this.pseudonymizer.restore(textToRestore, mappings);
+
+    // Only modify if something changed
+    if (restoredText === textToRestore) {
+      return Promise.resolve(this.allow(context));
+    }
+
+    this.addModification(
+      context,
+      phase,
+      phase === 'post_response' ? 'response' : 'toolResult',
+      `Restored ${mappings.size} PII placeholders`
+    );
+
+    if (phase === 'post_response') {
+      context.response = restoredText;
+    } else {
+      try {
+        context.toolResult = JSON.parse(restoredText) as unknown;
+      } catch {
+        context.toolResult = restoredText;
+      }
+    }
+
+    return Promise.resolve(this.modify(context));
+  }
+
+  /**
+   * Get detector for testing
+   */
+  getDetector(): PIIDetector {
+    return this.detector;
+  }
+
+  /**
+   * Get pseudonymizer for testing
+   */
+  getPseudonymizer(): Pseudonymizer {
+    return this.pseudonymizer;
+  }
+}

package/src/guardrails/plugins/pii-redactor/pseudonymizer.ts

@@ -0,0 +1,91 @@
+import { PIIDetection, PIIType } from './detectors.js';
+
+/**
+ * Pseudonymizer - replaces PII with numbered placeholders
+ * and supports optional restoration
+ */
+export class Pseudonymizer {
+  private counters: Map<PIIType, number> = new Map();
+
+  /**
+   * Pseudonymize text by replacing PII with placeholders
+   * Returns the modified text and a mapping for restoration
+   */
+  pseudonymize(
+    text: string,
+    detections: PIIDetection[]
+  ): { text: string; mappings: Map<string, string> } {
+    const mappings = new Map<string, string>();
+    let result = text;
+    let offset = 0;
+
+    // Reset counters for consistent numbering
+    this.counters.clear();
+
+    for (const detection of detections) {
+      const placeholder = this.generatePlaceholder(detection.type);
+      mappings.set(placeholder, detection.value);
+
+      // Replace in text (accounting for previous replacements)
+      const start = detection.startIndex + offset;
+      const end = detection.endIndex + offset;
+      result = result.substring(0, start) + placeholder + result.substring(end);
+
+      // Adjust offset for next replacement
+      offset += placeholder.length - detection.value.length;
+    }
+
+    return { text: result, mappings };
+  }
+
+  /**
+   * Restore placeholders in text with original values
+   */
+  restore(text: string, mappings: Map<string, string>): string {
+    let result = text;
+
+    for (const [placeholder, original] of mappings) {
+      // Use global replace to handle multiple occurrences
+      result = result.replace(
+        new RegExp(this.escapeRegex(placeholder), 'g'),
+        original
+      );
+    }
+
+    return result;
+  }
+
+  /**
+   * Generate a placeholder for a PII type
+   */
+  private generatePlaceholder(type: PIIType): string {
+    const count = (this.counters.get(type) || 0) + 1;
+    this.counters.set(type, count);
+
+    const typeLabels: Record<PIIType, string> = {
+      email: 'EMAIL',
+      phone: 'PHONE',
+      ssn: 'SSN',
+      api_key: 'API_KEY',
+      credit_card: 'CARD',
+      ip_address: 'IP',
+      custom: 'REDACTED',
+    };
+
+    return `[${typeLabels[type]}_${count}]`;
+  }
+
+  /**
+   * Escape special regex characters in a string
+   */
+  private escapeRegex(str: string): string {
+    return str.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
+  }
+
+  /**
+   * Reset counters (for testing)
+   */
+  reset(): void {
+    this.counters.clear();
+  }
+}

package/src/guardrails/plugins/rate-limiter.ts

@@ -0,0 +1,142 @@
+import { BaseGuardrailPlugin } from './base-plugin.js';
+import { GuardrailPhase, GuardrailContext, GuardrailResult } from '../types.js';
+import { RateLimiterConfig } from '../../config/types.js';
+
+interface RequestRecord {
+  timestamp: number;
+}
+
+/**
+ * Rate limiter plugin - limits requests per minute/hour
+ */
+export class RateLimiterPlugin extends BaseGuardrailPlugin {
+  name = 'rate_limiter';
+  phases: GuardrailPhase[] = ['pre_request'];
+
+  private requestsPerMinute: number = 60;
+  private requestsPerHour: number = 1000;
+  private perProvider: boolean = false;
+  private burstAllowance: number = 5;
+
+  // Request history: key is provider (or 'global'), value is array of timestamps
+  private requestHistory: Map<string, RequestRecord[]> = new Map();
+
+  async initialize(config: Record<string, unknown>): Promise<void> {
+    await super.initialize(config);
+
+    const typedConfig = config as Partial<RateLimiterConfig>;
+    this.requestsPerMinute = typedConfig.requests_per_minute ?? 60;
+    this.requestsPerHour = typedConfig.requests_per_hour ?? 1000;
+    this.perProvider = typedConfig.per_provider ?? false;
+    this.burstAllowance = typedConfig.burst_allowance ?? 5;
+    this.priority = typedConfig.priority ?? 10;
+  }
+
+  execute(phase: GuardrailPhase, context: GuardrailContext): Promise<GuardrailResult> {
+    if (phase !== 'pre_request') {
+      return Promise.resolve(this.allow(context));
+    }
+
+    const key = this.perProvider ? context.provider : 'global';
+    const now = Date.now();
+
+    // Get or create request history for this key
+    let history = this.requestHistory.get(key);
+    if (!history) {
+      history = [];
+      this.requestHistory.set(key, history);
+    }
+
+    // Clean up old entries (older than 1 hour)
+    const oneHourAgo = now - 60 * 60 * 1000;
+    history = history.filter((r) => r.timestamp > oneHourAgo);
+
+    // Remove empty keys to prevent unbounded Map growth with perProvider mode
+    if (history.length === 0) {
+      this.requestHistory.delete(key);
+      history = [];
+    } else {
+      this.requestHistory.set(key, history);
+    }
+
+    // Count requests in last minute and last hour
+    const oneMinuteAgo = now - 60 * 1000;
+    const requestsLastMinute = history.filter((r) => r.timestamp > oneMinuteAgo).length;
+    const requestsLastHour = history.length;
+
+    // Check rate limits (with burst allowance)
+    const effectiveMinuteLimit = this.requestsPerMinute + this.burstAllowance;
+    const effectiveHourLimit = this.requestsPerHour + this.burstAllowance;
+
+    if (requestsLastMinute >= effectiveMinuteLimit) {
+      this.addViolation(
+        context,
+        phase,
+        'requests_per_minute',
+        'error',
+        `Rate limit exceeded: ${requestsLastMinute} requests in the last minute (limit: ${this.requestsPerMinute})`,
+        { requestsLastMinute, limit: this.requestsPerMinute }
+      );
+      return Promise.resolve(this.block(
+        context,
+        `Rate limit exceeded: ${requestsLastMinute}/${this.requestsPerMinute} requests per minute`
+      ));
+    }
+
+    if (requestsLastHour >= effectiveHourLimit) {
+      this.addViolation(
+        context,
+        phase,
+        'requests_per_hour',
+        'error',
+        `Rate limit exceeded: ${requestsLastHour} requests in the last hour (limit: ${this.requestsPerHour})`,
+        { requestsLastHour, limit: this.requestsPerHour }
+      );
+      return Promise.resolve(this.block(
+        context,
+        `Rate limit exceeded: ${requestsLastHour}/${this.requestsPerHour} requests per hour`
+      ));
+    }
+
+    // Log warning if approaching limit
+    if (requestsLastMinute >= this.requestsPerMinute * 0.8) {
+      this.addViolation(
+        context,
+        phase,
+        'requests_per_minute_warning',
+        'warning',
+        `Approaching rate limit: ${requestsLastMinute}/${this.requestsPerMinute} requests per minute`,
+        { requestsLastMinute, limit: this.requestsPerMinute }
+      );
+    }
+
+    // Record this request
+    history.push({ timestamp: now });
+    // Ensure history is stored in Map (needed after empty cleanup)
+    this.requestHistory.set(key, history);
+
+    return Promise.resolve(this.allow(context));
+  }
+
+  /**
+   * Get current request counts (for testing/monitoring)
+   */
+  getRequestCounts(key: string = 'global'): { lastMinute: number; lastHour: number } {
+    const now = Date.now();
+    const history = this.requestHistory.get(key) || [];
+    const oneMinuteAgo = now - 60 * 1000;
+    const oneHourAgo = now - 60 * 60 * 1000;
+
+    return {
+      lastMinute: history.filter((r) => r.timestamp > oneMinuteAgo).length,
+      lastHour: history.filter((r) => r.timestamp > oneHourAgo).length,
+    };
+  }
+
+  /**
+   * Reset request history (for testing)
+   */
+  reset(): void {
+    this.requestHistory.clear();
+  }
+}

package/src/guardrails/plugins/token-limiter.ts

@@ -0,0 +1,155 @@
+import { BaseGuardrailPlugin } from './base-plugin.js';
+import { GuardrailPhase, GuardrailContext, GuardrailResult } from '../types.js';
+import { TokenLimiterConfig } from '../../config/types.js';
+
+/**
+ * Token limiter plugin - limits input/output token counts
+ */
+export class TokenLimiterPlugin extends BaseGuardrailPlugin {
+  name = 'token_limiter';
+  phases: GuardrailPhase[] = ['pre_request', 'post_response'];
+
+  private maxInputTokens: number = 8192;
+  private maxOutputTokens: number | undefined;
+  private warnAtPercentage: number = 80;
+
+  async initialize(config: Record<string, unknown>): Promise<void> {
+    await super.initialize(config);
+
+    const typedConfig = config as Partial<TokenLimiterConfig>;
+    this.maxInputTokens = typedConfig.max_input_tokens ?? 8192;
+    this.maxOutputTokens = typedConfig.max_output_tokens;
+    this.warnAtPercentage = typedConfig.warn_at_percentage ?? 80;
+    this.priority = typedConfig.priority ?? 20;
+  }
+
+  execute(phase: GuardrailPhase, context: GuardrailContext): Promise<GuardrailResult> {
+    if (phase === 'pre_request') {
+      return this.checkInputTokens(context, phase);
+    } else if (phase === 'post_response') {
+      return this.checkOutputTokens(context, phase);
+    }
+    return Promise.resolve(this.allow(context));
+  }
+
+  private checkInputTokens(
+    context: GuardrailContext,
+    phase: GuardrailPhase
+  ): Promise<GuardrailResult> {
+    // Estimate token count from prompt
+    const prompt = context.prompt || '';
+    const estimatedTokens = this.estimateTokenCount(prompt);
+
+    // Also count messages if present
+    let totalTokens = estimatedTokens;
+    for (const msg of context.messages) {
+      totalTokens += this.estimateTokenCount(msg.content);
+    }
+
+    // Check if over limit
+    if (totalTokens > this.maxInputTokens) {
+      this.addViolation(
+        context,
+        phase,
+        'max_input_tokens',
+        'error',
+        `Token limit exceeded: estimated ${totalTokens} tokens (limit: ${this.maxInputTokens})`,
+        { estimatedTokens: totalTokens, limit: this.maxInputTokens }
+      );
+      return Promise.resolve(
+        this.block(context, `Token limit exceeded: ~${totalTokens}/${this.maxInputTokens} tokens`)
+      );
+    }
+
+    // Warn if approaching limit
+    const warnThreshold = this.maxInputTokens * (this.warnAtPercentage / 100);
+    if (totalTokens >= warnThreshold) {
+      this.addViolation(
+        context,
+        phase,
+        'max_input_tokens_warning',
+        'warning',
+        `Approaching token limit: estimated ${totalTokens}/${this.maxInputTokens} tokens (${Math.round((totalTokens / this.maxInputTokens) * 100)}%)`,
+        {
+          estimatedTokens: totalTokens,
+          limit: this.maxInputTokens,
+          percentage: Math.round((totalTokens / this.maxInputTokens) * 100),
+        }
+      );
+    }
+
+    return Promise.resolve(this.allow(context));
+  }
+
+  private checkOutputTokens(
+    context: GuardrailContext,
+    phase: GuardrailPhase
+  ): Promise<GuardrailResult> {
+    // Skip if no output limit configured
+    if (!this.maxOutputTokens) {
+      return Promise.resolve(this.allow(context));
+    }
+
+    const response = context.response || '';
+    const estimatedTokens = this.estimateTokenCount(response);
+
+    // Check if over limit
+    if (estimatedTokens > this.maxOutputTokens) {
+      this.addViolation(
+        context,
+        phase,
+        'max_output_tokens',
+        'error',
+        `Output token limit exceeded: estimated ${estimatedTokens} tokens (limit: ${this.maxOutputTokens})`,
+        { estimatedTokens, limit: this.maxOutputTokens }
+      );
+      return Promise.resolve(
+        this.block(
+          context,
+          `Output token limit exceeded: ~${estimatedTokens}/${this.maxOutputTokens} tokens`
+        )
+      );
+    }
+
+    // Warn if approaching limit
+    const warnThreshold = this.maxOutputTokens * (this.warnAtPercentage / 100);
+    if (estimatedTokens >= warnThreshold) {
+      this.addViolation(
+        context,
+        phase,
+        'max_output_tokens_warning',
+        'warning',
+        `Approaching output token limit: estimated ${estimatedTokens}/${this.maxOutputTokens} tokens (${Math.round((estimatedTokens / this.maxOutputTokens) * 100)}%)`,
+        {
+          estimatedTokens,
+          limit: this.maxOutputTokens,
+          percentage: Math.round((estimatedTokens / this.maxOutputTokens) * 100),
+        }
+      );
+    }
+
+    return Promise.resolve(this.allow(context));
+  }
+
+  /**
+   * Estimate token count from text
+   * Uses a simple heuristic: ~4 characters per token for English text
+   * This is a rough approximation - for accuracy, use tiktoken
+   */
+  estimateTokenCount(text: string): number {
+    if (!text) return 0;
+    // Rough approximation: 1 token ≈ 4 characters for English
+    // Add some overhead for special tokens
+    return Math.ceil(text.length / 4) + 4;
+  }
+
+  /**
+   * Get configured limits (for testing/monitoring)
+   */
+  getLimits(): { maxInputTokens: number; maxOutputTokens: number | undefined } {
+    return {
+      maxInputTokens: this.maxInputTokens,
+      maxOutputTokens: this.maxOutputTokens,
+    };
+  }
+}