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

import { describe, it, expect, vi } from 'vitest';
import { LlmPromptSelector, type PromptIndexForLlm } from '../src/gate/llm-selector.js';
import { ProviderRegistry } from '../src/providers/registry.js';
import type { LlmProvider, CompletionOptions, CompletionResult } from '../src/providers/types.js';

function makeMockProvider(responseContent: string): LlmProvider {
  return {
    name: 'mock-heavy',
    complete: vi.fn().mockResolvedValue({
      content: responseContent,
      toolCalls: [],
      usage: { promptTokens: 100, completionTokens: 50, totalTokens: 150 },
      finishReason: 'stop',
    } satisfies CompletionResult),
    listModels: vi.fn().mockResolvedValue(['mock-model']),
    isAvailable: vi.fn().mockResolvedValue(true),
  };
}

function makeRegistry(provider: LlmProvider): ProviderRegistry {
  const registry = new ProviderRegistry();
  registry.register(provider);
  registry.assignTier(provider.name, 'heavy');
  return registry;
}

const sampleIndex: PromptIndexForLlm[] = [
  { name: 'zigbee-pairing', priority: 7, summary: 'How to pair Zigbee devices', chapters: ['Setup', 'Troubleshooting'] },
  { name: 'mqtt-config', priority: 5, summary: 'MQTT broker configuration', chapters: null },
  { name: 'common-mistakes', priority: 10, summary: 'Critical safety rules', chapters: null },
];

describe('LlmPromptSelector', () => {
  it('sends tags and index to heavy LLM and parses response', async () => {
    const provider = makeMockProvider(
      '```json\n{ "selectedNames": ["zigbee-pairing"], "reasoning": "User is working with zigbee" }\n```',
    );
    const registry = makeRegistry(provider);
    const selector = new LlmPromptSelector(registry);

    const result = await selector.selectPrompts(['zigbee', 'pairing'], sampleIndex);

    expect(result.selectedNames).toContain('zigbee-pairing');
    expect(result.selectedNames).toContain('common-mistakes'); // Priority 10 always included
    expect(result.reasoning).toBe('User is working with zigbee');
  });

  it('always includes priority 10 prompts even if LLM omits them', async () => {
    const provider = makeMockProvider(
      '{ "selectedNames": ["mqtt-config"], "reasoning": "MQTT related" }',
    );
    const registry = makeRegistry(provider);
    const selector = new LlmPromptSelector(registry);

    const result = await selector.selectPrompts(['mqtt'], sampleIndex);

    expect(result.selectedNames).toContain('mqtt-config');
    expect(result.selectedNames).toContain('common-mistakes');
  });

  it('does not duplicate priority 10 if LLM already selected them', async () => {
    const provider = makeMockProvider(
      '{ "selectedNames": ["common-mistakes", "mqtt-config"], "reasoning": "Both needed" }',
    );
    const registry = makeRegistry(provider);
    const selector = new LlmPromptSelector(registry);

    const result = await selector.selectPrompts(['mqtt'], sampleIndex);

    const count = result.selectedNames.filter((n) => n === 'common-mistakes').length;
    expect(count).toBe(1);
  });

  it('passes system and user messages to provider.complete', async () => {
    const provider = makeMockProvider(
      '{ "selectedNames": [], "reasoning": "none" }',
    );
    const registry = makeRegistry(provider);
    const selector = new LlmPromptSelector(registry);

    await selector.selectPrompts(['test'], sampleIndex);

    expect(provider.complete).toHaveBeenCalledOnce();
    const call = (provider.complete as ReturnType<typeof vi.fn>).mock.calls[0]![0] as CompletionOptions;
    expect(call.messages).toHaveLength(2);
    expect(call.messages[0]!.role).toBe('system');
    expect(call.messages[1]!.role).toBe('user');
    expect(call.messages[1]!.content).toContain('test');
    expect(call.temperature).toBe(0);
  });

  it('passes model override to complete options', async () => {
    const provider = makeMockProvider(
      '{ "selectedNames": [], "reasoning": "" }',
    );
    const registry = makeRegistry(provider);
    const selector = new LlmPromptSelector(registry, 'gemini-pro');

    await selector.selectPrompts(['test'], sampleIndex);

    const call = (provider.complete as ReturnType<typeof vi.fn>).mock.calls[0]![0] as CompletionOptions;
    expect(call.model).toBe('gemini-pro');
  });

  it('throws when no heavy provider is available', async () => {
    const registry = new ProviderRegistry(); // Empty registry
    const selector = new LlmPromptSelector(registry);

    await expect(selector.selectPrompts(['test'], sampleIndex)).rejects.toThrow(
      'No heavy LLM provider available',
    );
  });

  it('throws when LLM response has no valid JSON', async () => {
    const provider = makeMockProvider('I cannot help with that request.');
    const registry = makeRegistry(provider);
    const selector = new LlmPromptSelector(registry);

    await expect(selector.selectPrompts(['test'], sampleIndex)).rejects.toThrow(
      'LLM response did not contain valid selection JSON',
    );
  });

  it('handles response with empty selectedNames', async () => {
    const provider = makeMockProvider('{ "selectedNames": [], "reasoning": "nothing matched" }');
    const registry = makeRegistry(provider);
    const selector = new LlmPromptSelector(registry);

    // Empty selectedNames, but priority 10 should still be included
    const result = await selector.selectPrompts(['test'], sampleIndex);
    expect(result.selectedNames).toEqual(['common-mistakes']);
    expect(result.reasoning).toBe('nothing matched');
  });

  it('handles response with reasoning missing', async () => {
    const provider = makeMockProvider('{ "selectedNames": ["mqtt-config"] }');
    const registry = makeRegistry(provider);
    const selector = new LlmPromptSelector(registry);

    const result = await selector.selectPrompts(['test'], sampleIndex);
    expect(result.reasoning).toBe('');
    expect(result.selectedNames).toContain('mqtt-config');
  });

  it('includes prompt details in the user prompt', async () => {
    const indexWithNull: PromptIndexForLlm[] = [
      ...sampleIndex,
      { name: 'no-desc', priority: 3, summary: null, chapters: null },
    ];
    const provider = makeMockProvider(
      '{ "selectedNames": [], "reasoning": "" }',
    );
    const registry = makeRegistry(provider);
    const selector = new LlmPromptSelector(registry);

    await selector.selectPrompts(['zigbee'], indexWithNull);

    const call = (provider.complete as ReturnType<typeof vi.fn>).mock.calls[0]![0] as CompletionOptions;
    const userMsg = call.messages[1]!.content;
    expect(userMsg).toContain('zigbee-pairing');
    expect(userMsg).toContain('priority: 7');
    expect(userMsg).toContain('How to pair Zigbee devices');
    expect(userMsg).toContain('Setup, Troubleshooting');
    expect(userMsg).toContain('No summary'); // For prompts with null summary
  });
});