Momento/memento-note/tests/unit/router.test.ts

import { describe, it, expect, afterEach } from 'vitest';
import { resolveAiRoute, resolveAiRouteWithTiming } from '@/lib/ai/router';
import { getProviderInstance, type ProviderType } from '@/lib/ai/factory';

describe('resolveAiRoute', () => {
  afterEach(() => {
    delete process.env.AI_PROVIDER_TAGS;
    delete process.env.AI_PROVIDER_EMBEDDING;
    delete process.env.AI_PROVIDER;
    delete process.env.AI_PROVIDER_CHAT;
    delete process.env.AI_MODEL_TAGS;
    delete process.env.AI_MODEL_EMBEDDING;
    delete process.env.AI_MODEL_CHAT;
    delete process.env.LMSTUDIO_API_KEY;
  });

  it('tags lane prefers AI_PROVIDER_TAGS over embedding/generic fallbacks', () => {
    const route = resolveAiRoute('tags', {
      AI_PROVIDER_TAGS: 'openrouter',
      AI_PROVIDER_EMBEDDING: 'deepseek',
      AI_PROVIDER: 'openai',
      AI_MODEL_TAGS: 'anthropic/claude-3.5-haiku',
      AI_MODEL_EMBEDDING: 'openai/text-embedding-3-small',
    });
    expect(route.lane).toBe('tags');
    expect(route.providerType).toBe('openrouter');
    expect(route.modelName).toBe('anthropic/claude-3.5-haiku');
    expect(route.embeddingModelName).toBe('openai/text-embedding-3-small');
  });

  it('embedding lane rejects anthropic gateways', () => {
    expect(() =>
      resolveAiRoute('embedding', {
        AI_PROVIDER_EMBEDDING: 'anthropic',
        AI_MODEL_TAGS: 'x',
        AI_MODEL_EMBEDDING: 'y',
      })
    ).toThrow(/AI_PROVIDER_EMBEDDING cannot use/);
  });

  it('chat openrouter lane keeps slash-model IDs', () => {
    const route = resolveAiRoute('chat', {
      AI_PROVIDER_CHAT: 'openrouter',
      AI_MODEL_CHAT: 'deepseek/deepseek-v4-flash',
      AI_MODEL_EMBEDDING: 'openai/text-embedding-3-small',
    });
    expect(route.providerType).toBe('openrouter');
    expect(route.modelName).toBe('deepseek/deepseek-v4-flash');
  });

  it('throws on unknown provider string', () => {
    expect(() =>
      resolveAiRoute('chat', {
        AI_PROVIDER_CHAT: 'unknown_provider',
      })
    ).toThrow(/Unknown AI provider 'unknown_provider'/);
  });

  it('throws on provider with typo', () => {
    expect(() =>
      resolveAiRoute('chat', {
        AI_PROVIDER_CHAT: 'open ai',
      })
    ).toThrow(/Unknown AI provider 'open ai'/);
  });

  it('uses provider-specific default for deepseek when no model configured', () => {
    const route = resolveAiRoute('chat', {
      AI_PROVIDER_CHAT: 'deepseek',
    });
    expect(route.providerType).toBe('deepseek');
    expect(route.modelName).toBe('deepseek-chat');
    expect(route.embeddingModelName).toBe('');
  });

  it('uses provider-specific default for openai when no model configured', () => {
    const route = resolveAiRoute('chat', {
      AI_PROVIDER_CHAT: 'openai',
    });
    expect(route.providerType).toBe('openai');
    expect(route.modelName).toBe('gpt-4o-mini');
    expect(route.embeddingModelName).toBe('text-embedding-3-small');
  });

  it('uses provider-specific default for google when no model configured', () => {
    const route = resolveAiRoute('chat', {
      AI_PROVIDER_CHAT: 'google',
    });
    expect(route.providerType).toBe('google');
    expect(route.modelName).toBe('gemini-1.5-flash');
    expect(route.embeddingModelName).toBe('text-embedding-004');
  });

  it('uses provider-specific default for mistral when no model configured', () => {
    const route = resolveAiRoute('chat', {
      AI_PROVIDER_CHAT: 'mistral',
    });
    expect(route.providerType).toBe('mistral');
    expect(route.modelName).toBe('mistral-small-latest');
    expect(route.embeddingModelName).toBe('mistral-embed');
  });

  it('uses provider-specific default for openrouter when no model configured', () => {
    const route = resolveAiRoute('chat', {
      AI_PROVIDER_CHAT: 'openrouter',
    });
    expect(route.providerType).toBe('openrouter');
    expect(route.modelName).toBe('openai/gpt-4o-mini');
    expect(route.embeddingModelName).toBe('openai/text-embedding-3-small');
  });

  it('explicit model overrides provider default', () => {
    const route = resolveAiRoute('chat', {
      AI_PROVIDER_CHAT: 'openai',
      AI_MODEL_CHAT: 'gpt-5.4',
    });
    expect(route.modelName).toBe('gpt-5.4');
  });

  it('handles null config values without crash', () => {
    const route = resolveAiRoute('chat', {
      AI_PROVIDER_CHAT: 'ollama',
      AI_MODEL_CHAT: null as unknown as string,
    });
    expect(route.providerType).toBe('ollama');
    expect(route.modelName).toBe('granite4:latest');
  });

  it('resolveAiRoute median stays under 50ms (warm CPU path)', () => {
    const cfg = {
      AI_PROVIDER_CHAT: 'deepseek',
      AI_MODEL_CHAT: 'deepseek-v4-flash',
      AI_MODEL_EMBEDDING: '',
    };
    const iterations = 400;
    const samples: number[] = [];
    for (let i = 0; i < iterations; i++) {
      const t0 = performance.now();
      resolveAiRoute('chat', cfg);
      samples.push(performance.now() - t0);
    }
    samples.sort((a, b) => a - b);
    const median = samples[Math.floor(samples.length / 2)]!;
    expect(median).toBeLessThan(50);
  });

  it('resolve + instantiate (lmstudio) median stays under 50ms on warm path', () => {
    process.env.LMSTUDIO_API_KEY = 'lm-studio';
    const cfg: Record<string, string> = {
      AI_PROVIDER_CHAT: 'lmstudio',
      AI_MODEL_CHAT: 'local-model',
      AI_MODEL_EMBEDDING: 'local-embed',
    };
    const iterations = 100;
    const samples: number[] = [];
    for (let i = 0; i < iterations; i++) {
      const route = resolveAiRoute('chat', cfg);
      const t0 = performance.now();
      getProviderInstance(route.providerType as ProviderType, cfg, route.modelName, route.embeddingModelName, route.ollamaBaseUrl);
      samples.push(performance.now() - t0);
    }
    samples.sort((a, b) => a - b);
    const median = samples[Math.floor(samples.length / 2)]!;
    expect(median).toBeLessThan(50);
  });

  it('resolveAiRouteWithTiming injects meta.resolveMs', () => {
    const route = resolveAiRouteWithTiming('chat', {
      AI_PROVIDER_CHAT: 'openai',
      AI_MODEL_CHAT: 'gpt-4.1-mini',
      AI_MODEL_EMBEDDING: 'text-embedding-3-small',
    });
    expect(route.meta.resolveMs).toBeDefined();
    expect(route.meta.resolveMs!).toBeGreaterThanOrEqual(0);
    expect(route.meta.resolveMs!).toBeLessThan(50);
  });
});