mcp-rubber-duck 1.8.0 → 1.9.0

package/tests/guardrails/plugins/rate-limiter.test.ts

@@ -0,0 +1,310 @@
+import { describe, it, expect, beforeEach, jest } from '@jest/globals';
+import { RateLimiterPlugin } from '../../../src/guardrails/plugins/rate-limiter';
+import { createGuardrailContext } from '../../../src/guardrails/context';
+
+// Mock logger to avoid console noise during tests
+jest.mock('../../../src/utils/logger');
+
+describe('RateLimiterPlugin', () => {
+  let plugin: RateLimiterPlugin;
+
+  beforeEach(async () => {
+    plugin = new RateLimiterPlugin();
+    plugin.reset();
+  });
+
+  describe('initialization', () => {
+    it('should initialize with default values', async () => {
+      await plugin.initialize({ enabled: true });
+
+      expect(plugin.enabled).toBe(true);
+      expect(plugin.name).toBe('rate_limiter');
+      expect(plugin.phases).toContain('pre_request');
+    });
+
+    it('should initialize with custom config', async () => {
+      await plugin.initialize({
+        enabled: true,
+        priority: 5,
+        requests_per_minute: 30,
+        requests_per_hour: 500,
+        per_provider: true,
+        burst_allowance: 10,
+      });
+
+      expect(plugin.priority).toBe(5);
+    });
+  });
+
+  describe('rate limiting - per minute', () => {
+    it('should allow requests under the limit', async () => {
+      await plugin.initialize({
+        enabled: true,
+        requests_per_minute: 10,
+        requests_per_hour: 1000,
+        burst_allowance: 0,
+      });
+
+      const context = createGuardrailContext({ provider: 'openai' });
+
+      // Make 5 requests - all should be allowed
+      for (let i = 0; i < 5; i++) {
+        const result = await plugin.execute('pre_request', context);
+        expect(result.action).toBe('allow');
+      }
+
+      expect(plugin.getRequestCounts().lastMinute).toBe(5);
+    });
+
+    it('should block requests over the limit', async () => {
+      await plugin.initialize({
+        enabled: true,
+        requests_per_minute: 3,
+        requests_per_hour: 1000,
+        burst_allowance: 0,
+      });
+
+      const context = createGuardrailContext({ provider: 'openai' });
+
+      // Make 3 requests - all should be allowed
+      for (let i = 0; i < 3; i++) {
+        const result = await plugin.execute('pre_request', context);
+        expect(result.action).toBe('allow');
+      }
+
+      // 4th request should be blocked
+      const result = await plugin.execute('pre_request', context);
+      expect(result.action).toBe('block');
+      expect(result.blockedBy).toBe('rate_limiter');
+      expect(result.blockReason).toContain('Rate limit exceeded');
+    });
+
+    it('should allow burst requests within burst allowance', async () => {
+      await plugin.initialize({
+        enabled: true,
+        requests_per_minute: 3,
+        requests_per_hour: 1000,
+        burst_allowance: 2,
+      });
+
+      const context = createGuardrailContext({ provider: 'openai' });
+
+      // Make 5 requests - all should be allowed (3 + 2 burst)
+      for (let i = 0; i < 5; i++) {
+        const result = await plugin.execute('pre_request', context);
+        expect(result.action).toBe('allow');
+      }
+
+      // 6th request should be blocked
+      const result = await plugin.execute('pre_request', context);
+      expect(result.action).toBe('block');
+    });
+  });
+
+  describe('rate limiting - per provider', () => {
+    it('should track requests globally by default', async () => {
+      await plugin.initialize({
+        enabled: true,
+        requests_per_minute: 5,
+        requests_per_hour: 1000,
+        per_provider: false,
+        burst_allowance: 0,
+      });
+
+      // Make requests for different providers
+      for (let i = 0; i < 5; i++) {
+        const provider = i % 2 === 0 ? 'openai' : 'gemini';
+        const context = createGuardrailContext({ provider });
+        await plugin.execute('pre_request', context);
+      }
+
+      // Global count should be 5
+      expect(plugin.getRequestCounts('global').lastMinute).toBe(5);
+
+      // Next request from any provider should be blocked
+      const context = createGuardrailContext({ provider: 'openai' });
+      const result = await plugin.execute('pre_request', context);
+      expect(result.action).toBe('block');
+    });
+
+    it('should track requests per provider when configured', async () => {
+      await plugin.initialize({
+        enabled: true,
+        requests_per_minute: 3,
+        requests_per_hour: 1000,
+        per_provider: true,
+        burst_allowance: 0,
+      });
+
+      // Make 3 requests for OpenAI
+      for (let i = 0; i < 3; i++) {
+        const context = createGuardrailContext({ provider: 'openai' });
+        const result = await plugin.execute('pre_request', context);
+        expect(result.action).toBe('allow');
+      }
+
+      // OpenAI should be blocked
+      const openaiContext = createGuardrailContext({ provider: 'openai' });
+      const openaiResult = await plugin.execute('pre_request', openaiContext);
+      expect(openaiResult.action).toBe('block');
+
+      // But Gemini should still be allowed
+      const geminiContext = createGuardrailContext({ provider: 'gemini' });
+      const geminiResult = await plugin.execute('pre_request', geminiContext);
+      expect(geminiResult.action).toBe('allow');
+
+      expect(plugin.getRequestCounts('openai').lastMinute).toBe(3);
+      expect(plugin.getRequestCounts('gemini').lastMinute).toBe(1);
+    });
+  });
+
+  describe('rate limiting - per hour', () => {
+    it('should block requests when hourly limit exceeded', async () => {
+      await plugin.initialize({
+        enabled: true,
+        requests_per_minute: 1000, // High minute limit
+        requests_per_hour: 5, // Low hour limit
+        burst_allowance: 0,
+      });
+
+      const context = createGuardrailContext({ provider: 'openai' });
+
+      // Make 5 requests - all should be allowed
+      for (let i = 0; i < 5; i++) {
+        const result = await plugin.execute('pre_request', context);
+        expect(result.action).toBe('allow');
+      }
+
+      // 6th request should be blocked due to hourly limit
+      const result = await plugin.execute('pre_request', context);
+      expect(result.action).toBe('block');
+      expect(result.blockReason).toContain('per hour');
+    });
+
+    it('should add violation for hourly limit exceeded', async () => {
+      await plugin.initialize({
+        enabled: true,
+        requests_per_minute: 1000,
+        requests_per_hour: 3,
+        burst_allowance: 0,
+      });
+
+      const context = createGuardrailContext({ provider: 'openai' });
+
+      // Make 3 requests
+      for (let i = 0; i < 3; i++) {
+        await plugin.execute('pre_request', context);
+      }
+
+      // 4th request - should be blocked with hourly violation
+      const result = await plugin.execute('pre_request', context);
+
+      expect(result.action).toBe('block');
+      const hourlyViolation = context.violations.find((v) => v.rule === 'requests_per_hour');
+      expect(hourlyViolation).toBeDefined();
+      expect(hourlyViolation?.severity).toBe('error');
+    });
+  });
+
+  describe('violations', () => {
+    it('should add violation when rate limit exceeded', async () => {
+      await plugin.initialize({
+        enabled: true,
+        requests_per_minute: 1,
+        requests_per_hour: 1000,
+        burst_allowance: 0,
+      });
+
+      const context = createGuardrailContext({ provider: 'openai' });
+
+      // First request
+      await plugin.execute('pre_request', context);
+
+      // Second request - should be blocked with violation
+      const result = await plugin.execute('pre_request', context);
+
+      expect(result.action).toBe('block');
+      expect(context.violations.length).toBeGreaterThan(0);
+      expect(context.violations[0].rule).toBe('requests_per_minute');
+      expect(context.violations[0].severity).toBe('error');
+    });
+
+    it('should add warning when approaching limit', async () => {
+      await plugin.initialize({
+        enabled: true,
+        requests_per_minute: 10,
+        requests_per_hour: 1000,
+        burst_allowance: 0,
+      });
+
+      const context = createGuardrailContext({ provider: 'openai' });
+
+      // Make 9 requests - the 9th sees 8 previous requests (80% of limit)
+      // Warning is added when request count before current >= 80%
+      for (let i = 0; i < 9; i++) {
+        await plugin.execute('pre_request', context);
+      }
+
+      // Check for warning violation
+      const warnings = context.violations.filter(
+        (v) => v.severity === 'warning' && v.rule === 'requests_per_minute_warning'
+      );
+      expect(warnings.length).toBeGreaterThan(0);
+    });
+  });
+
+  describe('phase handling', () => {
+    it('should only process pre_request phase', async () => {
+      await plugin.initialize({ enabled: true, requests_per_minute: 1 });
+
+      const context = createGuardrailContext({ provider: 'openai' });
+
+      // pre_request should be processed
+      const preResult = await plugin.execute('pre_request', context);
+      expect(preResult.action).toBe('allow');
+
+      // post_response should be skipped (not in phases)
+      const postResult = await plugin.execute('post_response', context);
+      expect(postResult.action).toBe('allow');
+
+      // Only 1 request counted (pre_request)
+      expect(plugin.getRequestCounts().lastMinute).toBe(1);
+    });
+  });
+
+  describe('shutdown', () => {
+    it('should disable plugin on shutdown', async () => {
+      await plugin.initialize({ enabled: true, requests_per_minute: 10 });
+
+      expect(plugin.enabled).toBe(true);
+
+      await plugin.shutdown();
+
+      expect(plugin.enabled).toBe(false);
+    });
+  });
+
+  describe('reset', () => {
+    it('should clear request history on reset', async () => {
+      await plugin.initialize({
+        enabled: true,
+        requests_per_minute: 10,
+        burst_allowance: 0,
+      });
+
+      const context = createGuardrailContext({ provider: 'openai' });
+
+      // Make some requests
+      for (let i = 0; i < 5; i++) {
+        await plugin.execute('pre_request', context);
+      }
+
+      expect(plugin.getRequestCounts().lastMinute).toBe(5);
+
+      // Reset
+      plugin.reset();
+
+      expect(plugin.getRequestCounts().lastMinute).toBe(0);
+    });
+  });
+});

package/tests/guardrails/plugins/token-limiter.test.ts

@@ -0,0 +1,216 @@
+import { describe, it, expect, beforeEach, jest } from '@jest/globals';
+import { TokenLimiterPlugin } from '../../../src/guardrails/plugins/token-limiter';
+import { createGuardrailContext } from '../../../src/guardrails/context';
+
+// Mock logger to avoid console noise during tests
+jest.mock('../../../src/utils/logger');
+
+describe('TokenLimiterPlugin', () => {
+  let plugin: TokenLimiterPlugin;
+
+  beforeEach(async () => {
+    plugin = new TokenLimiterPlugin();
+  });
+
+  describe('initialization', () => {
+    it('should initialize with default values', async () => {
+      await plugin.initialize({ enabled: true });
+
+      expect(plugin.enabled).toBe(true);
+      expect(plugin.name).toBe('token_limiter');
+      expect(plugin.phases).toContain('pre_request');
+      expect(plugin.getLimits().maxInputTokens).toBe(8192);
+    });
+
+    it('should initialize with custom config', async () => {
+      await plugin.initialize({
+        enabled: true,
+        priority: 15,
+        max_input_tokens: 4096,
+        max_output_tokens: 2048,
+        warn_at_percentage: 90,
+      });
+
+      expect(plugin.priority).toBe(15);
+      const limits = plugin.getLimits();
+      expect(limits.maxInputTokens).toBe(4096);
+      expect(limits.maxOutputTokens).toBe(2048);
+    });
+  });
+
+  describe('token estimation', () => {
+    it('should estimate tokens from text', async () => {
+      await plugin.initialize({ enabled: true });
+
+      // ~4 chars per token
+      expect(plugin.estimateTokenCount('Hello')).toBeGreaterThan(0);
+      expect(plugin.estimateTokenCount('Hello world')).toBeGreaterThan(plugin.estimateTokenCount('Hello'));
+    });
+
+    it('should return 0 for empty text', async () => {
+      await plugin.initialize({ enabled: true });
+
+      expect(plugin.estimateTokenCount('')).toBe(0);
+    });
+  });
+
+  describe('token limiting', () => {
+    it('should allow prompts under the limit', async () => {
+      await plugin.initialize({
+        enabled: true,
+        max_input_tokens: 1000,
+      });
+
+      const shortPrompt = 'Hello world'; // ~7 tokens
+      const context = createGuardrailContext({ prompt: shortPrompt });
+      const result = await plugin.execute('pre_request', context);
+
+      expect(result.action).toBe('allow');
+    });
+
+    it('should block prompts over the limit', async () => {
+      await plugin.initialize({
+        enabled: true,
+        max_input_tokens: 10, // Very small limit
+      });
+
+      // Create a prompt that exceeds the limit (~4 chars per token)
+      const longPrompt = 'This is a longer prompt that will definitely exceed our tiny token limit';
+      const context = createGuardrailContext({ prompt: longPrompt });
+      const result = await plugin.execute('pre_request', context);
+
+      expect(result.action).toBe('block');
+      expect(result.blockedBy).toBe('token_limiter');
+      expect(result.blockReason).toContain('Token limit exceeded');
+    });
+
+    it('should include messages in token count', async () => {
+      await plugin.initialize({
+        enabled: true,
+        max_input_tokens: 20,
+      });
+
+      const context = createGuardrailContext({
+        prompt: 'Short',
+        messages: [
+          { role: 'user', content: 'This is a message with many tokens that should push us over the limit', timestamp: new Date() },
+        ],
+      });
+
+      const result = await plugin.execute('pre_request', context);
+      expect(result.action).toBe('block');
+    });
+  });
+
+  describe('warnings', () => {
+    it('should add warning when approaching limit', async () => {
+      await plugin.initialize({
+        enabled: true,
+        max_input_tokens: 100,
+        warn_at_percentage: 50, // Warn at 50%
+      });
+
+      // Create prompt that's ~60-80% of limit
+      const prompt = 'A'.repeat(300); // ~75 tokens
+      const context = createGuardrailContext({ prompt });
+      await plugin.execute('pre_request', context);
+
+      const warnings = context.violations.filter((v) => v.severity === 'warning');
+      expect(warnings.length).toBeGreaterThan(0);
+    });
+
+    it('should not warn when well under limit', async () => {
+      await plugin.initialize({
+        enabled: true,
+        max_input_tokens: 1000,
+        warn_at_percentage: 80,
+      });
+
+      const shortPrompt = 'Hello';
+      const context = createGuardrailContext({ prompt: shortPrompt });
+      await plugin.execute('pre_request', context);
+
+      const warnings = context.violations.filter((v) => v.severity === 'warning');
+      expect(warnings.length).toBe(0);
+    });
+  });
+
+  describe('phase handling', () => {
+    it('should process pre_request and post_response phases', async () => {
+      await plugin.initialize({ enabled: true });
+
+      expect(plugin.phases).toContain('pre_request');
+      expect(plugin.phases).toContain('post_response');
+    });
+
+    it('should allow other phases', async () => {
+      await plugin.initialize({ enabled: true });
+
+      const context = createGuardrailContext({ prompt: 'test' });
+
+      // pre_tool_input should pass through
+      const result = await plugin.execute('pre_tool_input', context);
+      expect(result.action).toBe('allow');
+    });
+  });
+
+  describe('output token limiting', () => {
+    it('should allow responses under the limit', async () => {
+      await plugin.initialize({
+        enabled: true,
+        max_output_tokens: 1000,
+      });
+
+      const context = createGuardrailContext({ prompt: 'test' });
+      context.response = 'Short response';
+      const result = await plugin.execute('post_response', context);
+
+      expect(result.action).toBe('allow');
+    });
+
+    it('should block responses over the limit', async () => {
+      await plugin.initialize({
+        enabled: true,
+        max_output_tokens: 10, // Very small limit
+      });
+
+      const context = createGuardrailContext({ prompt: 'test' });
+      context.response = 'This is a longer response that will definitely exceed our tiny token limit for output';
+      const result = await plugin.execute('post_response', context);
+
+      expect(result.action).toBe('block');
+      expect(result.blockedBy).toBe('token_limiter');
+      expect(result.blockReason).toContain('Output token limit exceeded');
+    });
+
+    it('should skip output check if no max_output_tokens configured', async () => {
+      await plugin.initialize({
+        enabled: true,
+        // No max_output_tokens
+      });
+
+      const context = createGuardrailContext({ prompt: 'test' });
+      context.response = 'A'.repeat(10000); // Very long response
+      const result = await plugin.execute('post_response', context);
+
+      expect(result.action).toBe('allow');
+    });
+
+    it('should add warning when response approaches output limit', async () => {
+      await plugin.initialize({
+        enabled: true,
+        max_output_tokens: 100,
+        warn_at_percentage: 50, // Warn at 50%
+      });
+
+      // Create response that's ~60-80% of limit
+      const context = createGuardrailContext({ prompt: 'test' });
+      context.response = 'A'.repeat(300); // ~75 tokens
+      await plugin.execute('post_response', context);
+
+      const warnings = context.violations.filter((v) => v.severity === 'warning');
+      expect(warnings.length).toBeGreaterThan(0);
+      expect(warnings[0].rule).toBe('max_output_tokens_warning');
+    });
+  });
+});