mcpctl/src/mcplocal/tests/proxymodel-llm-adapter.test.ts

import { describe, it, expect, vi } from 'vitest';
import { LLMProviderAdapter } from '../src/proxymodel/llm-adapter.js';
import { ProviderRegistry } from '../src/providers/registry.js';
import type { LlmProvider, CompletionResult } from '../src/providers/types.js';

function mockProvider(name: string, response = 'mock response'): LlmProvider {
  return {
    name,
    complete: vi.fn().mockResolvedValue({
      content: response,
      toolCalls: [],
      usage: { promptTokens: 10, completionTokens: 5, totalTokens: 15 },
      finishReason: 'stop',
    } satisfies CompletionResult),
    listModels: vi.fn().mockResolvedValue([]),
    isAvailable: vi.fn().mockResolvedValue(true),
  };
}

function failingProvider(name: string, error = 'connection refused'): LlmProvider {
  return {
    name,
    complete: vi.fn().mockRejectedValue(new Error(error)),
    listModels: vi.fn().mockResolvedValue([]),
    isAvailable: vi.fn().mockResolvedValue(true),
  };
}

describe('LLMProviderAdapter', () => {
  it('available() returns true when a provider is registered', () => {
    const registry = new ProviderRegistry();
    registry.register(mockProvider('test'));
    registry.assignTier('test', 'fast');

    const adapter = new LLMProviderAdapter(registry);
    expect(adapter.available()).toBe(true);
  });

  it('available() returns false when no provider is registered', () => {
    const registry = new ProviderRegistry();
    const adapter = new LLMProviderAdapter(registry);
    expect(adapter.available()).toBe(false);
  });

  it('complete() sends prompt as user message', async () => {
    const provider = mockProvider('test');
    const registry = new ProviderRegistry();
    registry.register(provider);
    registry.assignTier('test', 'fast');

    const adapter = new LLMProviderAdapter(registry);
    const result = await adapter.complete('summarize this');

    expect(result).toBe('mock response');
    expect(provider.complete).toHaveBeenCalledWith({
      messages: [{ role: 'user', content: 'summarize this' }],
      temperature: 0,
    });
  });

  it('complete() includes system message when provided', async () => {
    const provider = mockProvider('test');
    const registry = new ProviderRegistry();
    registry.register(provider);
    registry.assignTier('test', 'fast');

    const adapter = new LLMProviderAdapter(registry);
    await adapter.complete('summarize', { system: 'You are a summarizer', maxTokens: 200 });

    expect(provider.complete).toHaveBeenCalledWith({
      messages: [
        { role: 'system', content: 'You are a summarizer' },
        { role: 'user', content: 'summarize' },
      ],
      maxTokens: 200,
      temperature: 0,
    });
  });

  it('complete() throws when no provider available', async () => {
    const registry = new ProviderRegistry();
    const adapter = new LLMProviderAdapter(registry);

    await expect(adapter.complete('test')).rejects.toThrow('No LLM provider available');
  });

  // --- Failover tests ---

  it('falls back to next provider in same tier on failure', async () => {
    const failing = failingProvider('vllm-local', 'vLLM startup timed out');
    const working = mockProvider('anthropic', 'anthropic response');
    const registry = new ProviderRegistry();
    registry.register(failing);
    registry.register(working);
    registry.assignTier('vllm-local', 'fast');
    registry.assignTier('anthropic', 'fast');

    const adapter = new LLMProviderAdapter(registry);
    const result = await adapter.complete('test');

    expect(result).toBe('anthropic response');
    expect(failing.complete).toHaveBeenCalledOnce();
    expect(working.complete).toHaveBeenCalledOnce();
  });

  it('falls back cross-tier when all fast providers fail', async () => {
    const fastFail = failingProvider('vllm-local');
    const heavy = mockProvider('gemini', 'gemini response');
    const registry = new ProviderRegistry();
    registry.register(fastFail);
    registry.register(heavy);
    registry.assignTier('vllm-local', 'fast');
    registry.assignTier('gemini', 'heavy');

    const adapter = new LLMProviderAdapter(registry);
    const result = await adapter.complete('test');

    expect(result).toBe('gemini response');
    expect(fastFail.complete).toHaveBeenCalledOnce();
    expect(heavy.complete).toHaveBeenCalledOnce();
  });

  it('throws last error when all providers fail', async () => {
    const fail1 = failingProvider('vllm', 'vLLM down');
    const fail2 = failingProvider('anthropic', 'rate limited');
    const registry = new ProviderRegistry();
    registry.register(fail1);
    registry.register(fail2);
    registry.assignTier('vllm', 'fast');
    registry.assignTier('anthropic', 'heavy');

    const adapter = new LLMProviderAdapter(registry);
    await expect(adapter.complete('test')).rejects.toThrow('rate limited');
    expect(fail1.complete).toHaveBeenCalledOnce();
    expect(fail2.complete).toHaveBeenCalledOnce();
  });

  it('does not retry provider that already succeeded', async () => {
    const fast = mockProvider('fast-provider', 'fast result');
    const heavy = mockProvider('heavy-provider', 'heavy result');
    const registry = new ProviderRegistry();
    registry.register(fast);
    registry.register(heavy);
    registry.assignTier('fast-provider', 'fast');
    registry.assignTier('heavy-provider', 'heavy');

    const adapter = new LLMProviderAdapter(registry);
    const result = await adapter.complete('test');

    expect(result).toBe('fast result');
    expect(fast.complete).toHaveBeenCalledOnce();
    expect(heavy.complete).not.toHaveBeenCalled();
  });

  it('prefers named provider but falls back on failure', async () => {
    const named = failingProvider('preferred', 'preferred down');
    const fallback = mockProvider('fallback', 'fallback response');
    const registry = new ProviderRegistry();
    registry.register(named);
    registry.register(fallback);
    registry.assignTier('preferred', 'fast');
    registry.assignTier('fallback', 'fast');

    const adapter = new LLMProviderAdapter(registry, 'preferred');
    const result = await adapter.complete('test');

    expect(result).toBe('fallback response');
  });
});