mcpctl/src/mcplocal/tests/llm-processor.test.ts

import { describe, it, expect, vi } from 'vitest';
import { LlmProcessor, DEFAULT_PROCESSOR_CONFIG } from '../src/llm/processor.js';
import { ProviderRegistry } from '../src/providers/registry.js';
import type { LlmProvider, CompletionResult } from '../src/providers/types.js';

function mockProvider(responses: string[]): LlmProvider {
  let callIndex = 0;
  return {
    name: 'mock',
    async complete(): Promise<CompletionResult> {
      const content = responses[callIndex] ?? '{}';
      callIndex++;
      return {
        content,
        toolCalls: [],
        usage: { promptTokens: 10, completionTokens: 5, totalTokens: 15 },
        finishReason: 'stop',
      };
    },
    async listModels() { return ['mock-1']; },
    async isAvailable() { return true; },
  };
}

function makeRegistry(provider?: LlmProvider): ProviderRegistry {
  const registry = new ProviderRegistry();
  if (provider) {
    registry.register(provider);
  }
  return registry;
}

describe('LlmProcessor.shouldProcess', () => {
  it('bypasses protocol-level methods', () => {
    const proc = new LlmProcessor(makeRegistry());
    expect(proc.shouldProcess('initialize')).toBe(false);
    expect(proc.shouldProcess('tools/list')).toBe(false);
    expect(proc.shouldProcess('resources/list')).toBe(false);
    expect(proc.shouldProcess('prompts/list')).toBe(false);
  });

  it('returns false when no tool name', () => {
    const proc = new LlmProcessor(makeRegistry());
    expect(proc.shouldProcess('tools/call')).toBe(false);
  });

  it('returns true for normal tool calls', () => {
    const proc = new LlmProcessor(makeRegistry());
    expect(proc.shouldProcess('tools/call', 'slack/search_messages')).toBe(true);
  });

  it('skips excluded tools', () => {
    const proc = new LlmProcessor(makeRegistry(), {
      ...DEFAULT_PROCESSOR_CONFIG,
      excludeTools: ['slack'],
    });
    expect(proc.shouldProcess('tools/call', 'slack/search_messages')).toBe(false);
    expect(proc.shouldProcess('tools/call', 'github/search')).toBe(true);
  });

  it('skips simple CRUD operations', () => {
    const proc = new LlmProcessor(makeRegistry());
    expect(proc.shouldProcess('tools/call', 'slack/create_channel')).toBe(false);
    expect(proc.shouldProcess('tools/call', 'slack/delete_message')).toBe(false);
    expect(proc.shouldProcess('tools/call', 'slack/remove_user')).toBe(false);
  });
});

describe('LlmProcessor.preprocessRequest', () => {
  it('returns original params when preprocessing disabled', async () => {
    const proc = new LlmProcessor(makeRegistry(mockProvider(['{}'])), {
      ...DEFAULT_PROCESSOR_CONFIG,
      enablePreprocessing: false,
    });
    const result = await proc.preprocessRequest('slack/search', { query: 'test' });
    expect(result.optimized).toBe(false);
    expect(result.params).toEqual({ query: 'test' });
  });

  it('returns original params when no provider', async () => {
    const proc = new LlmProcessor(makeRegistry(), {
      ...DEFAULT_PROCESSOR_CONFIG,
      enablePreprocessing: true,
    });
    const result = await proc.preprocessRequest('slack/search', { query: 'test' });
    expect(result.optimized).toBe(false);
  });

  it('optimizes params with LLM', async () => {
    const provider = mockProvider([JSON.stringify({ query: 'test', limit: 10 })]);
    const proc = new LlmProcessor(makeRegistry(provider), {
      ...DEFAULT_PROCESSOR_CONFIG,
      enablePreprocessing: true,
    });
    const result = await proc.preprocessRequest('slack/search', { query: 'test' });
    expect(result.optimized).toBe(true);
    expect(result.params).toEqual({ query: 'test', limit: 10 });
  });

  it('falls back on LLM error', async () => {
    const badProvider: LlmProvider = {
      name: 'bad',
      async complete() { throw new Error('LLM down'); },
      async listModels() { return []; },
      async isAvailable() { return false; },
    };
    const proc = new LlmProcessor(makeRegistry(badProvider), {
      ...DEFAULT_PROCESSOR_CONFIG,
      enablePreprocessing: true,
    });
    const result = await proc.preprocessRequest('slack/search', { query: 'test' });
    expect(result.optimized).toBe(false);
    expect(result.params).toEqual({ query: 'test' });
  });
});

describe('LlmProcessor.filterResponse', () => {
  it('returns original when filtering disabled', async () => {
    const proc = new LlmProcessor(makeRegistry(mockProvider([])), {
      ...DEFAULT_PROCESSOR_CONFIG,
      enableFiltering: false,
    });
    const response = { jsonrpc: '2.0' as const, id: '1', result: { data: 'big' } };
    const result = await proc.filterResponse('slack/search', response);
    expect(result.filtered).toBe(false);
  });

  it('returns original when no provider', async () => {
    const proc = new LlmProcessor(makeRegistry());
    const response = { jsonrpc: '2.0' as const, id: '1', result: { data: 'x'.repeat(600) } };
    const result = await proc.filterResponse('slack/search', response);
    expect(result.filtered).toBe(false);
  });

  it('skips small responses below token threshold', async () => {
    const proc = new LlmProcessor(makeRegistry(mockProvider([])));
    // With default tokenThreshold=250, any response < 1000 chars (~250 tokens) is skipped
    const response = { jsonrpc: '2.0' as const, id: '1', result: { data: 'small' } };
    const result = await proc.filterResponse('slack/search', response);
    expect(result.filtered).toBe(false);
  });

  it('skips error responses', async () => {
    const proc = new LlmProcessor(makeRegistry(mockProvider([])));
    const response = { jsonrpc: '2.0' as const, id: '1', error: { code: -1, message: 'fail' } };
    const result = await proc.filterResponse('slack/search', response);
    expect(result.filtered).toBe(false);
  });

  it('filters large responses with LLM', async () => {
    const largeData = { items: Array.from({ length: 50 }, (_, i) => ({ id: i, name: `item-${i}`, extra: 'x'.repeat(20) })) };
    const filteredData = { items: [{ id: 0, name: 'item-0' }, { id: 1, name: 'item-1' }] };
    const provider = mockProvider([JSON.stringify(filteredData)]);

    const proc = new LlmProcessor(makeRegistry(provider));
    const response = { jsonrpc: '2.0' as const, id: '1', result: largeData };
    const result = await proc.filterResponse('slack/search', response);
    expect(result.filtered).toBe(true);
    expect(result.filteredSize).toBeLessThan(result.originalSize);
  });

  it('falls back on LLM error', async () => {
    const badProvider: LlmProvider = {
      name: 'bad',
      async complete() { throw new Error('LLM down'); },
      async listModels() { return []; },
      async isAvailable() { return false; },
    };
    const largeData = { items: Array.from({ length: 50 }, (_, i) => ({ id: i, extra: 'x'.repeat(20) })) };
    const proc = new LlmProcessor(makeRegistry(badProvider));
    const response = { jsonrpc: '2.0' as const, id: '1', result: largeData };
    const result = await proc.filterResponse('slack/search', response);
    expect(result.filtered).toBe(false);
    expect(result.result).toEqual(largeData);
  });

  it('respects custom tokenThreshold', async () => {
    // Set a very high threshold so that even "big" responses are skipped
    const proc = new LlmProcessor(makeRegistry(mockProvider([])), {
      ...DEFAULT_PROCESSOR_CONFIG,
      tokenThreshold: 10_000,
    });
    const largeData = { items: Array.from({ length: 50 }, (_, i) => ({ id: i, name: `item-${i}` })) };
    const response = { jsonrpc: '2.0' as const, id: '1', result: largeData };
    const result = await proc.filterResponse('slack/search', response);
    expect(result.filtered).toBe(false);
  });

  it('uses filter cache to skip repeated filtering', async () => {
    // First call: LLM returns same-size data => cache records shouldFilter=false
    const largeData = { items: Array.from({ length: 50 }, (_, i) => ({ id: i, extra: 'x'.repeat(20) })) };
    const raw = JSON.stringify(largeData);
    // Return something larger so the cache stores shouldFilter=false (filtered not smaller)
    const notSmaller = JSON.stringify(largeData);
    const provider = mockProvider([notSmaller]);

    const proc = new LlmProcessor(makeRegistry(provider));
    const response = { jsonrpc: '2.0' as const, id: '1', result: largeData };

    // First call goes to LLM
    await proc.filterResponse('slack/search', response);

    // Second call should hit cache (shouldFilter=false) and skip LLM
    const result2 = await proc.filterResponse('slack/search', response);
    expect(result2.filtered).toBe(false);

    const metrics = proc.getMetrics();
    expect(metrics.cacheHits).toBeGreaterThanOrEqual(1);
  });

  it('records metrics on filter operations', async () => {
    const largeData = { items: Array.from({ length: 50 }, (_, i) => ({ id: i, name: `item-${i}`, extra: 'x'.repeat(20) })) };
    const filteredData = { items: [{ id: 0, name: 'item-0' }] };
    const provider = mockProvider([JSON.stringify(filteredData)]);

    const proc = new LlmProcessor(makeRegistry(provider));
    const response = { jsonrpc: '2.0' as const, id: '1', result: largeData };
    await proc.filterResponse('slack/search', response);

    const metrics = proc.getMetrics();
    expect(metrics.filterCount).toBe(1);
    expect(metrics.totalTokensProcessed).toBeGreaterThan(0);
    expect(metrics.tokensSaved).toBeGreaterThan(0);
    expect(metrics.cacheMisses).toBe(1);
  });

  it('records metrics even on LLM failure', async () => {
    const badProvider: LlmProvider = {
      name: 'bad',
      async complete() { throw new Error('LLM down'); },
      async listModels() { return []; },
      async isAvailable() { return false; },
    };
    const largeData = { items: Array.from({ length: 50 }, (_, i) => ({ id: i, extra: 'x'.repeat(20) })) };
    const proc = new LlmProcessor(makeRegistry(badProvider));
    const response = { jsonrpc: '2.0' as const, id: '1', result: largeData };
    await proc.filterResponse('slack/search', response);

    const metrics = proc.getMetrics();
    expect(metrics.filterCount).toBe(1);
    expect(metrics.totalTokensProcessed).toBeGreaterThan(0);
    // No tokens saved because filter failed
    expect(metrics.tokensSaved).toBe(0);
  });
});

describe('LlmProcessor metrics and cache management', () => {
  it('exposes metrics via getMetrics()', () => {
    const proc = new LlmProcessor(makeRegistry());
    const metrics = proc.getMetrics();
    expect(metrics.totalTokensProcessed).toBe(0);
    expect(metrics.filterCount).toBe(0);
  });

  it('resets metrics', async () => {
    const largeData = { items: Array.from({ length: 50 }, (_, i) => ({ id: i, extra: 'x'.repeat(20) })) };
    const provider = mockProvider([JSON.stringify({ summary: 'ok' })]);
    const proc = new LlmProcessor(makeRegistry(provider));
    const response = { jsonrpc: '2.0' as const, id: '1', result: largeData };
    await proc.filterResponse('slack/search', response);

    expect(proc.getMetrics().filterCount).toBe(1);
    proc.resetMetrics();
    expect(proc.getMetrics().filterCount).toBe(0);
  });

  it('clears filter cache', async () => {
    const largeData = { items: Array.from({ length: 50 }, (_, i) => ({ id: i, extra: 'x'.repeat(20) })) };
    const filteredData = { items: [{ id: 0 }] };
    // Two responses needed: first call filters, second call after cache clear also filters
    const provider = mockProvider([JSON.stringify(filteredData), JSON.stringify(filteredData)]);
    const proc = new LlmProcessor(makeRegistry(provider));
    const response = { jsonrpc: '2.0' as const, id: '1', result: largeData };

    await proc.filterResponse('slack/search', response);
    proc.clearFilterCache();

    // After clearing cache, should get a cache miss again
    proc.resetMetrics();
    await proc.filterResponse('slack/search', response);
    expect(proc.getMetrics().cacheMisses).toBe(1);
  });
});