stack/apps/gateway/src/agent/__tests__/routing.service.test.ts

import { describe, it, expect, beforeEach, vi } from 'vitest';
import { RoutingService } from '../routing.service.js';
import type { ModelInfo } from '@mosaic/types';

const mockModels: ModelInfo[] = [
  {
    id: 'claude-3-haiku',
    provider: 'anthropic',
    name: 'Claude 3 Haiku',
    reasoning: false,
    contextWindow: 200_000,
    maxTokens: 4096,
    inputTypes: ['text', 'image'],
    cost: { input: 0.25, output: 1.25, cacheRead: 0.03, cacheWrite: 0.3 },
  },
  {
    id: 'claude-3-sonnet',
    provider: 'anthropic',
    name: 'Claude 3 Sonnet',
    reasoning: true,
    contextWindow: 200_000,
    maxTokens: 8192,
    inputTypes: ['text', 'image'],
    cost: { input: 3, output: 15, cacheRead: 0.3, cacheWrite: 3.75 },
  },
  {
    id: 'llama3.2',
    provider: 'ollama',
    name: 'Llama 3.2',
    reasoning: false,
    contextWindow: 128_000,
    maxTokens: 4096,
    inputTypes: ['text'],
    cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
  },
];

function createMockProviderService() {
  return {
    listAvailableModels: vi.fn().mockReturnValue(mockModels),
    findModel: vi.fn(),
    getDefaultModel: vi.fn(),
    getRegistry: vi.fn(),
    listProviders: vi.fn(),
    registerCustomProvider: vi.fn(),
  };
}

describe('RoutingService', () => {
  let routingService: RoutingService;
  let mockProviderService: ReturnType<typeof createMockProviderService>;

  beforeEach(() => {
    mockProviderService = createMockProviderService();
    routingService = new RoutingService(mockProviderService as never);
  });

  it('returns a model when no criteria specified', () => {
    const result = routingService.route();
    expect(result).not.toBeNull();
    expect(result!.provider).toBeDefined();
    expect(result!.modelId).toBeDefined();
    expect(result!.score).toBeGreaterThan(0);
  });

  it('returns null when no models available', () => {
    mockProviderService.listAvailableModels.mockReturnValue([]);
    const result = routingService.route();
    expect(result).toBeNull();
  });

  it('selects preferred model when specified', () => {
    const result = routingService.route({
      preferredProvider: 'anthropic',
      preferredModel: 'claude-3-sonnet',
    });
    expect(result).not.toBeNull();
    expect(result!.provider).toBe('anthropic');
    expect(result!.modelId).toBe('claude-3-sonnet');
    expect(result!.score).toBe(100);
  });

  it('disqualifies models without reasoning when required', () => {
    const result = routingService.route({ requireReasoning: true });
    expect(result).not.toBeNull();
    expect(result!.modelId).toBe('claude-3-sonnet');
  });

  it('disqualifies models without image input when required', () => {
    const result = routingService.route({ requireImageInput: true });
    expect(result).not.toBeNull();
    // Llama doesn't support images, should be excluded
    expect(result!.provider).toBe('anthropic');
  });

  it('respects minimum context window', () => {
    const result = routingService.route({ minContextWindow: 150_000 });
    expect(result).not.toBeNull();
    // Only anthropic models have 200k context
    expect(result!.provider).toBe('anthropic');
  });

  it('favors cheap models when costTier is cheap', () => {
    const result = routingService.route({ costTier: 'cheap' });
    expect(result).not.toBeNull();
    // Ollama (free) and Haiku ($0.25/M) are cheap
    expect(['ollama', 'anthropic'].includes(result!.provider)).toBe(true);
    if (result!.provider === 'anthropic') {
      expect(result!.modelId).toBe('claude-3-haiku');
    }
  });

  it('ranks all models and returns sorted results', () => {
    const ranked = routingService.rank({ taskType: 'coding' });
    expect(ranked.length).toBeGreaterThan(0);
    // Should be sorted by score descending
    for (let i = 1; i < ranked.length; i++) {
      expect(ranked[i]!.score).toBeLessThanOrEqual(ranked[i - 1]!.score);
    }
  });

  it('gives reasoning bonus for coding tasks', () => {
    const ranked = routingService.rank({ taskType: 'coding' });
    const sonnet = ranked.find((r) => r.modelId === 'claude-3-sonnet');
    const haiku = ranked.find((r) => r.modelId === 'claude-3-haiku');
    expect(sonnet).toBeDefined();
    expect(haiku).toBeDefined();
    // Sonnet (reasoning) should score higher for coding than haiku (no reasoning)
    expect(sonnet!.score).toBeGreaterThan(haiku!.score);
  });

  it('prefers specified provider', () => {
    const ranked = routingService.rank({ preferredProvider: 'ollama' });
    const ollamaModel = ranked.find((r) => r.provider === 'ollama');
    expect(ollamaModel).toBeDefined();
    expect(ollamaModel!.reasoning).toContain('preferred provider');
  });
});