proffesor-for-testing__agentic-qe/tests/unit/mcp/security/sampling-server.test.ts

/**
 * Agentic QE v3 - Sampling Server Tests
 * Tests for server-initiated LLM sampling
 */

import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
import {
  SamplingServer,
  createSamplingServer,
  QEDecisionPrompts,
} from '../../../../src/mcp/security/sampling-server';

describe('SamplingServer', () => {
  let server: SamplingServer;

  beforeEach(() => {
    server = createSamplingServer({
      maxRequestsPerMinute: 60,
      maxTokensPerMinute: 100000,
      enableCaching: true,
      cacheTTL: 5000,
    });
  });

  afterEach(() => {
    server.dispose();
  });

  describe('request creation', () => {
    it('should create a sampling request', () => {
      const request = server.createRequest({
        systemPrompt: 'You are a helpful assistant',
        messages: [
          { role: 'user', content: 'Hello' },
        ],
      });

      expect(request.requestId).toBeDefined();
      expect(request.requestId).toMatch(/^sr-/);
      expect(request.messages).toHaveLength(1);
    });

    it('should include optional parameters', () => {
      const request = server.createRequest({
        systemPrompt: 'System prompt',
        messages: [{ role: 'user', content: 'Test' }],
        model: 'claude-3-opus',
        maxTokens: 2000,
        temperature: 0.5,
      });

      expect(request.model).toBe('claude-3-opus');
      expect(request.maxTokens).toBe(2000);
      expect(request.temperature).toBe(0.5);
    });
  });

  describe('sampling', () => {
    it('should process sampling request', async () => {
      const request = server.createRequest({
        systemPrompt: 'You are a helpful assistant',
        messages: [
          { role: 'user', content: 'What is 2+2?' },
        ],
      });

      const response = await server.sample(request);

      expect(response.requestId).toBe(request.requestId);
      expect(response.content).toBeDefined();
      expect(response.stopReason).toBe('end_turn');
      expect(response.usage).toBeDefined();
      expect(response.usage.inputTokens).toBeGreaterThan(0);
      expect(response.usage.outputTokens).toBeGreaterThan(0);
    });

    it('should fill in default parameters', async () => {
      const request = server.createRequest({
        messages: [{ role: 'user', content: 'Test' }],
      });

      const response = await server.sample(request);

      expect(response.model).toBeDefined();
      expect(response.model).toContain('claude');
    });

    it('should handle rich content', async () => {
      const request = server.createRequest({
        messages: [
          {
            role: 'user',
            content: [
              { type: 'text', text: 'Analyze this code:' },
              { type: 'code', code: { language: 'typescript', content: 'const x = 1;' } },
            ],
          },
        ],
      });

      const response = await server.sample(request);
      expect(response.content).toBeDefined();
    });
  });

  describe('caching', () => {
    it('should cache identical requests', async () => {
      const request1 = server.createRequest({
        messages: [{ role: 'user', content: 'Cached test' }],
      });

      // First request - not cached
      const response1 = await server.sample(request1);
      expect(response1.cached).toBe(false);

      // Second request with same content - should be cached
      const request2 = server.createRequest({
        messages: [{ role: 'user', content: 'Cached test' }],
      });

      const response2 = await server.sample(request2);
      expect(response2.cached).toBe(true);
    });

    it('should clear cache', async () => {
      const request = server.createRequest({
        messages: [{ role: 'user', content: 'Clear cache test' }],
      });

      await server.sample(request);
      server.clearCache();

      const request2 = server.createRequest({
        messages: [{ role: 'user', content: 'Clear cache test' }],
      });

      const response = await server.sample(request2);
      expect(response.cached).toBe(false);
    });

    it('should expire cached items', async () => {
      vi.useFakeTimers();

      const request = server.createRequest({
        messages: [{ role: 'user', content: 'Expire test' }],
      });

      await server.sample(request);

      // Advance time past TTL
      vi.advanceTimersByTime(6000);

      const request2 = server.createRequest({
        messages: [{ role: 'user', content: 'Expire test' }],
      });

      const response = await server.sample(request2);
      expect(response.cached).toBe(false);

      vi.useRealTimers();
    });
  });

  describe('quota management', () => {
    it('should track request quota', async () => {
      const stats = server.getStats();
      expect(stats.remainingRequests).toBe(60);
    });

    it('should enforce request quota', async () => {
      // Create server with very low quota
      const limitedServer = createSamplingServer({
        maxRequestsPerMinute: 2,
        enableCaching: false, // Disable caching to test quota
      });

      // Make requests
      await limitedServer.sample(
        limitedServer.createRequest({ messages: [{ role: 'user', content: 'Request 1' }] })
      );
      await limitedServer.sample(
        limitedServer.createRequest({ messages: [{ role: 'user', content: 'Request 2' }] })
      );

      // Third request should fail
      await expect(
        limitedServer.sample(
          limitedServer.createRequest({ messages: [{ role: 'user', content: 'Request 3' }] })
        )
      ).rejects.toThrow(/Rate limit exceeded/);

      limitedServer.dispose();
    });

    it('should reset quota after window', async () => {
      vi.useFakeTimers();

      const limitedServer = createSamplingServer({
        maxRequestsPerMinute: 2,
        enableCaching: false,
      });

      // Exhaust quota
      await limitedServer.sample(
        limitedServer.createRequest({ messages: [{ role: 'user', content: 'R1' }] })
      );
      await limitedServer.sample(
        limitedServer.createRequest({ messages: [{ role: 'user', content: 'R2' }] })
      );

      // Advance time past 1 minute window
      vi.advanceTimersByTime(61000);

      // Should work now
      const response = await limitedServer.sample(
        limitedServer.createRequest({ messages: [{ role: 'user', content: 'R3' }] })
      );
      expect(response.content).toBeDefined();

      limitedServer.dispose();
      vi.useRealTimers();
    });
  });

  describe('custom handlers', () => {
    it('should register custom handler', async () => {
      server.registerHandler('custom', async (request) => ({
        requestId: request.requestId,
        content: 'Custom response',
        stopReason: 'end_turn',
        model: 'custom-model',
        usage: { inputTokens: 10, outputTokens: 5, totalTokens: 15 },
        cached: false,
        latencyMs: 0,
      }));

      const request = server.createRequest({
        messages: [{ role: 'user', content: 'Test' }],
      });

      const response = await server.sample(request, 'custom');

      expect(response.content).toBe('Custom response');
      expect(response.model).toBe('custom-model');
    });

    it('should throw for unknown handler', async () => {
      const request = server.createRequest({
        messages: [{ role: 'user', content: 'Test' }],
      });

      await expect(server.sample(request, 'nonexistent')).rejects.toThrow(
        /Unknown sampling handler/
      );
    });
  });

  describe('QE decision prompts', () => {
    it('should sample test generation decision', async () => {
      const response = await server.sampleQEDecision('testGenerationDecision', {
        sourceCode: 'function add(a: number, b: number) { return a + b; }',
        coverageData: { lineCoverage: 50, branchCoverage: 40 },
      });

      expect(response.content).toBeDefined();
      expect(response.content).toContain('Test');
    });

    it('should sample quality gate decision', async () => {
      const response = await server.sampleQEDecision('qualityGateDecision', {
        metrics: { coverage: 85, passRate: 99 },
        thresholds: { coverage: 80, passRate: 95 },
        trends: [
          { metric: 'coverage', direction: 'up', change: 5 },
        ],
      });

      expect(response.content).toBeDefined();
      expect(response.content).toContain('Quality Gate');
    });

    it('should sample defect prediction decision', async () => {
      const response = await server.sampleQEDecision('defectPredictionDecision', {
        codeChanges: [
          { file: 'auth.ts', additions: 100, deletions: 20, complexity: 15 },
        ],
        historicalDefects: [
          { file: 'auth.ts', defectCount: 5, lastDefect: '2024-01-01' },
        ],
        riskFactors: ['High complexity', 'Recent changes'],
      });

      expect(response.content).toBeDefined();
      expect(response.content).toContain('Defect');
    });

    it('should sample security analysis decision', async () => {
      const response = await server.sampleQEDecision('securityAnalysisDecision', {
        findings: [
          {
            type: 'SQL Injection',
            severity: 'high',
            location: 'user-service.ts:45',
            description: 'Unsanitized user input in query',
          },
        ],
        codeContext: 'db.query(`SELECT * FROM users WHERE id = ${userId}`)',
      });

      expect(response.content).toBeDefined();
    });
  });

  describe('statistics', () => {
    it('should track request statistics', async () => {
      await server.sample(
        server.createRequest({ messages: [{ role: 'user', content: 'Test 1' }] })
      );
      await server.sample(
        server.createRequest({ messages: [{ role: 'user', content: 'Test 2' }] })
      );

      const stats = server.getStats();

      expect(stats.totalRequests).toBe(2);
      expect(stats.totalTokensUsed).toBeGreaterThan(0);
      expect(stats.averageLatencyMs).toBeGreaterThanOrEqual(0);
    });

    it('should track cached requests', async () => {
      // First request
      await server.sample(
        server.createRequest({ messages: [{ role: 'user', content: 'Cache stats test' }] })
      );

      // Cached request
      await server.sample(
        server.createRequest({ messages: [{ role: 'user', content: 'Cache stats test' }] })
      );

      const stats = server.getStats();
      expect(stats.cachedRequests).toBe(1);
    });
  });

  describe('latency tracking', () => {
    it('should track latency', async () => {
      const response = await server.sample(
        server.createRequest({ messages: [{ role: 'user', content: 'Latency test' }] })
      );

      expect(response.latencyMs).toBeGreaterThanOrEqual(0);
    });
  });
});

describe('QEDecisionPrompts', () => {
  describe('testGenerationDecision', () => {
    it('should build proper prompt structure', () => {
      const prompt = QEDecisionPrompts.testGenerationDecision({
        sourceCode: 'function test() {}',
        existingTests: 'describe("test", () => {});',
        coverageData: { lineCoverage: 80, branchCoverage: 70 },
      });

      expect(prompt.systemPrompt).toContain('QE expert');
      expect(prompt.systemPrompt).toContain('test');
      expect(prompt.messages).toHaveLength(1);
      expect(prompt.messages[0].role).toBe('user');
    });

    it('should work without optional params', () => {
      const prompt = QEDecisionPrompts.testGenerationDecision({
        sourceCode: 'function test() {}',
      });

      expect(prompt.messages).toHaveLength(1);
    });
  });

  describe('qualityGateDecision', () => {
    it('should build proper prompt structure', () => {
      const prompt = QEDecisionPrompts.qualityGateDecision({
        metrics: { coverage: 85 },
        thresholds: { coverage: 80 },
        trends: [{ metric: 'coverage', direction: 'up', change: 5 }],
      });

      expect(prompt.systemPrompt).toContain('quality gate');
      expect(prompt.messages[0].content).toContain('Metrics');
    });
  });

  describe('defectPredictionDecision', () => {
    it('should build proper prompt structure', () => {
      const prompt = QEDecisionPrompts.defectPredictionDecision({
        codeChanges: [{ file: 'test.ts', additions: 10, deletions: 5, complexity: 3 }],
        historicalDefects: [{ file: 'test.ts', defectCount: 1, lastDefect: '2024-01-01' }],
        riskFactors: ['Complex code'],
      });

      expect(prompt.systemPrompt).toContain('defect');
      expect(prompt.messages[0].content).toContain('Code Changes');
    });
  });

  describe('securityAnalysisDecision', () => {
    it('should build proper prompt structure', () => {
      const prompt = QEDecisionPrompts.securityAnalysisDecision({
        findings: [
          { type: 'XSS', severity: 'high', location: 'file.ts:10', description: 'Unsafe innerHTML' },
        ],
        codeContext: 'element.innerHTML = userInput;',
      });

      expect(prompt.systemPrompt).toContain('security');
      expect(prompt.messages[0].content).toBeDefined();
    });
  });
});