mcp-rubber-duck 1.8.0 → 1.9.3

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,
+    };
+  }
+}

package/src/guardrails/service.ts

@@ -0,0 +1,209 @@
+import { GuardrailPlugin, GuardrailPhase, GuardrailContext, GuardrailResult, CreateContextOptions } from './types.js';
+import { createGuardrailContext } from './context.js';
+import { GuardrailsConfig, GuardrailsPluginsConfig } from '../config/types.js';
+import { logger } from '../utils/logger.js';
+
+/**
+ * Main service that orchestrates guardrail plugins
+ */
+export class GuardrailsService {
+  private plugins: GuardrailPlugin[] = [];
+  private config: GuardrailsConfig;
+  private enabled: boolean = false;
+
+  constructor(config?: Partial<GuardrailsConfig>) {
+    this.config = {
+      enabled: config?.enabled ?? false,
+      log_violations: config?.log_violations ?? true,
+      log_modifications: config?.log_modifications ?? false,
+      fail_open: config?.fail_open ?? false,
+      plugins: config?.plugins,
+    };
+    // Start disabled - will be enabled after successful initialization with plugins
+    this.enabled = false;
+  }
+
+  /**
+   * Initialize the service and all configured plugins
+   */
+  async initialize(): Promise<void> {
+    if (!this.config.enabled) {
+      logger.info('Guardrails disabled in configuration');
+      return;
+    }
+
+    const pluginConfigs = this.config.plugins || {};
+
+    // Load plugins in order
+    await this.loadPluginsFromConfig(pluginConfigs);
+
+    // Sort by priority (lower = runs first)
+    this.plugins.sort((a, b) => a.priority - b.priority);
+
+    this.enabled = this.plugins.length > 0;
+    logger.info(`Guardrails initialized with ${this.plugins.length} plugins`);
+  }
+
+  private async loadPluginsFromConfig(pluginConfigs: Partial<GuardrailsPluginsConfig>): Promise<void> {
+    const pluginOrder: Array<[string, unknown]> = [
+      ['rate_limiter', pluginConfigs.rate_limiter],
+      ['token_limiter', pluginConfigs.token_limiter],
+      ['pii_redactor', pluginConfigs.pii_redactor],
+      ['pattern_blocker', pluginConfigs.pattern_blocker],
+    ];
+
+    for (const [pluginName, pluginConfig] of pluginOrder) {
+      if (!pluginConfig || !(pluginConfig as { enabled?: boolean }).enabled) {
+        continue;
+      }
+
+      try {
+        const plugin = await this.loadPlugin(pluginName);
+        await plugin.initialize(pluginConfig as Record<string, unknown>);
+        if ((pluginConfig as { priority?: number }).priority !== undefined) {
+          plugin.priority = (pluginConfig as { priority: number }).priority;
+        }
+        this.plugins.push(plugin);
+        logger.info(`Guardrail plugin '${pluginName}' initialized`);
+      } catch (error) {
+        logger.error(`Failed to initialize guardrail plugin '${pluginName}':`, error);
+      }
+    }
+  }
+
+  private async loadPlugin(name: string): Promise<GuardrailPlugin> {
+    // Dynamic plugin loading
+    switch (name) {
+      case 'rate_limiter': {
+        const { RateLimiterPlugin } = await import('./plugins/rate-limiter.js');
+        return new RateLimiterPlugin();
+      }
+      case 'token_limiter': {
+        const { TokenLimiterPlugin } = await import('./plugins/token-limiter.js');
+        return new TokenLimiterPlugin();
+      }
+      case 'pattern_blocker': {
+        const { PatternBlockerPlugin } = await import('./plugins/pattern-blocker.js');
+        return new PatternBlockerPlugin();
+      }
+      case 'pii_redactor': {
+        const { PIIRedactorPlugin } = await import('./plugins/pii-redactor/index.js');
+        return new PIIRedactorPlugin();
+      }
+      default:
+        throw new Error(`Unknown guardrail plugin: ${name}`);
+    }
+  }
+
+  /**
+   * Check if guardrails are enabled
+   */
+  isEnabled(): boolean {
+    return this.enabled;
+  }
+
+  /**
+   * Create a new context for guardrail execution
+   */
+  createContext(options: CreateContextOptions): GuardrailContext {
+    return createGuardrailContext(options);
+  }
+
+  /**
+   * Execute all relevant plugins for a given phase
+   */
+  async execute(phase: GuardrailPhase, context: GuardrailContext): Promise<GuardrailResult> {
+    if (!this.enabled) {
+      return { action: 'allow', context };
+    }
+
+    const relevantPlugins = this.plugins.filter(
+      (p) => p.enabled && p.phases.includes(phase)
+    );
+
+    // Track logged items to avoid duplicates
+    let lastViolationCount = 0;
+    let lastModificationCount = 0;
+    let wasModified = false;
+
+    for (const plugin of relevantPlugins) {
+      try {
+        const result = await plugin.execute(phase, context);
+
+        // Log only NEW violations if configured
+        if (this.config.log_violations && context.violations.length > lastViolationCount) {
+          for (let i = lastViolationCount; i < context.violations.length; i++) {
+            const violation = context.violations[i];
+            logger.warn(`Guardrail violation: ${violation.pluginName} - ${violation.message}`, {
+              rule: violation.rule,
+              severity: violation.severity,
+              details: violation.details,
+            });
+          }
+          lastViolationCount = context.violations.length;
+        }
+
+        // Log only NEW modifications if configured
+        if (this.config.log_modifications && context.modifications.length > lastModificationCount) {
+          for (let i = lastModificationCount; i < context.modifications.length; i++) {
+            const mod = context.modifications[i];
+            logger.info(`Guardrail modification: ${mod.pluginName} - ${mod.reason}`, {
+              field: mod.field,
+            });
+          }
+          lastModificationCount = context.modifications.length;
+        }
+
+        if (result.action === 'block') {
+          logger.warn(`Request blocked by guardrail '${plugin.name}': ${result.blockReason}`);
+          return result;
+        }
+
+        // Track if any plugin modified the context
+        if (result.action === 'modify') {
+          wasModified = true;
+        }
+
+        // Update context for next plugin
+        context = result.context;
+      } catch (error) {
+        const errorMessage = error instanceof Error ? error.message : String(error);
+        logger.error(`Guardrail plugin '${plugin.name}' error:`, error);
+
+        if (!this.config.fail_open) {
+          return {
+            action: 'block',
+            context,
+            blockedBy: plugin.name,
+            blockReason: `Plugin error: ${errorMessage}`,
+          };
+        }
+        // fail_open: continue to next plugin
+      }
+    }
+
+    return { action: wasModified ? 'modify' : 'allow', context };
+  }
+
+  /**
+   * Shutdown the service and all plugins
+   */
+  async shutdown(): Promise<void> {
+    for (const plugin of this.plugins) {
+      try {
+        await plugin.shutdown();
+      } catch (error) {
+        logger.error(`Error shutting down plugin '${plugin.name}':`, error);
+      }
+    }
+    this.plugins = [];
+    this.enabled = false;
+  }
+
+  /**
+   * Get list of loaded plugins
+   */
+  getPlugins(): GuardrailPlugin[] {
+    return [...this.plugins];
+  }
+}

package/src/guardrails/types.ts

@@ -0,0 +1,120 @@
+import { ConversationMessage } from '../config/types.js';
+
+/**
+ * Phases in the guardrail pipeline where plugins can intercept
+ */
+export type GuardrailPhase =
+  | 'pre_request' // Before LLM API call
+  | 'post_response' // After LLM response, before tool handling
+  | 'pre_tool_input' // Before MCP tool execution
+  | 'post_tool_output' // After MCP tool returns
+  | 'pre_cache'; // Before caching response
+
+/**
+ * Action to take after guardrail evaluation
+ */
+export type GuardrailAction = 'allow' | 'block' | 'modify';
+
+/**
+ * Severity levels for violations
+ */
+export type ViolationSeverity = 'info' | 'warning' | 'error' | 'critical';
+
+/**
+ * A violation detected by a guardrail plugin
+ */
+export interface GuardrailViolation {
+  pluginName: string;
+  phase: GuardrailPhase;
+  rule: string;
+  severity: ViolationSeverity;
+  message: string;
+  details?: Record<string, unknown>;
+}
+
+/**
+ * A modification made by a guardrail plugin
+ */
+export interface GuardrailModification {
+  pluginName: string;
+  phase: GuardrailPhase;
+  field: string;
+  originalValue?: unknown;
+  newValue?: unknown;
+  reason: string;
+}
+
+/**
+ * Context passed through the guardrail pipeline
+ */
+export interface GuardrailContext {
+  // Request metadata
+  requestId: string;
+  provider: string;
+  model: string;
+  timestamp: Date;
+
+  // Phase-specific data (mutable by plugins)
+  messages: ConversationMessage[];
+  prompt?: string;
+  response?: string;
+  toolName?: string;
+  toolArgs?: Record<string, unknown>;
+  toolResult?: unknown;
+
+  // Tracking data (persisted across phases)
+  metadata: Map<string, unknown>; // For plugins to store state
+  violations: GuardrailViolation[]; // Accumulated violations
+  modifications: GuardrailModification[]; // Tracking changes made
+}
+
+/**
+ * Result from a guardrail plugin execution
+ */
+export interface GuardrailResult {
+  action: GuardrailAction;
+  context: GuardrailContext;
+  blockedBy?: string; // Plugin name that blocked
+  blockReason?: string;
+}
+
+/**
+ * Base interface for guardrail plugins
+ */
+export interface GuardrailPlugin {
+  /** Unique plugin name */
+  name: string;
+
+  /** Whether the plugin is currently enabled */
+  enabled: boolean;
+
+  /** Execution priority (lower = runs first) */
+  priority: number;
+
+  /** Which phases this plugin handles */
+  phases: GuardrailPhase[];
+
+  /** Initialize the plugin with its configuration */
+  initialize(config: Record<string, unknown>): Promise<void>;
+
+  /** Execute the plugin for a specific phase */
+  execute(phase: GuardrailPhase, context: GuardrailContext): Promise<GuardrailResult>;
+
+  /** Cleanup plugin resources */
+  shutdown(): Promise<void>;
+}
+
+/**
+ * Options for creating a guardrail context
+ */
+export interface CreateContextOptions {
+  requestId?: string;
+  provider?: string;
+  model?: string;
+  messages?: ConversationMessage[];
+  prompt?: string;
+  response?: string;
+  toolName?: string;
+  toolArgs?: Record<string, unknown>;
+  toolResult?: unknown;
+}