LexAI/tests/unit/providers.test.ts

import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import {
  callProvider,
  listModels,
  getSystemPrompt,
  defaultMaxTokens,
  providerLabel,
  PROVIDER_SPECS,
} from '@lib/providers';

function mockFetchOnce(data: unknown, { ok = true, status = 200 } = {}) {
  const fn = vi.fn().mockResolvedValue({ ok, status, json: async () => data });
  vi.stubGlobal('fetch', fn);
  return fn;
}

beforeEach(() => vi.restoreAllMocks());
afterEach(() => vi.unstubAllGlobals());

// ─── Request shapes pinned to the ORIGINAL callX/callXWithPrompt functions ────
// These bodies/headers/URLs are copied from the pre-refactor background.ts.
// maxTokens is forced to 1024 to match the old hard-coded value exactly.

const openAiResponse = { choices: [{ message: { content: '  fixed text  ' } }] };
const anthropicResponse = { content: [{ text: '  fixed text  ' }] };

describe('callProvider request shapes (parity with old implementations)', () => {
  it('OpenAI: url, bearer auth, system message, temperature 0.7, default model', async () => {
    const fetch = mockFetchOnce(openAiResponse);
    const res = await callProvider({ provider: 'openai', apiKey: 'sk-x' }, 'hello', 'SYS', { maxTokens: 1024 });

    expect(res).toEqual({ result: 'fixed text' });
    const [url, init] = fetch.mock.calls[0];
    expect(url).toBe('https://api.openai.com/v1/chat/completions');
    expect(init.method).toBe('POST');
    expect(init.headers).toEqual({ 'Content-Type': 'application/json', Authorization: 'Bearer sk-x' });
    expect(JSON.parse(init.body)).toEqual({
      model: 'gpt-4o-mini',
      messages: [
        { role: 'system', content: 'SYS' },
        { role: 'user', content: 'hello' },
      ],
      max_completion_tokens: 1024,
      temperature: 0.7,
    });
  });

  it('OpenAI reasoning models (o-series / gpt-5): max_completion_tokens, NO temperature', async () => {
    const fetch = mockFetchOnce(openAiResponse);
    await callProvider({ provider: 'openai', apiKey: 'sk-x', model: 'gpt-5-mini' }, 'hello', 'SYS', { maxTokens: 1024 });

    const body = JSON.parse(fetch.mock.calls[0][1].body);
    expect(body.max_completion_tokens).toBe(1024);
    expect(body.max_tokens).toBeUndefined();
    expect(body.temperature).toBeUndefined();
  });

  it('Anthropic: x-api-key + version headers, top-level system, no temperature', async () => {
    const fetch = mockFetchOnce(anthropicResponse);
    const res = await callProvider({ provider: 'anthropic', apiKey: 'sk-ant' }, 'hello', 'SYS', { maxTokens: 1024 });

    expect(res).toEqual({ result: 'fixed text' });
    const [url, init] = fetch.mock.calls[0];
    expect(url).toBe('https://api.anthropic.com/v1/messages');
    expect(init.headers).toEqual({
      'Content-Type': 'application/json',
      'x-api-key': 'sk-ant',
      'anthropic-version': '2023-06-01',
      'anthropic-dangerous-direct-browser-access': 'true',
    });
    expect(JSON.parse(init.body)).toEqual({
      model: 'claude-3-5-haiku-20241022',
      max_tokens: 1024,
      system: 'SYS',
      messages: [{ role: 'user', content: 'hello' }],
    });
  });

  it('Groq: OpenAI-compatible endpoint and body with temperature', async () => {
    const fetch = mockFetchOnce(openAiResponse);
    await callProvider({ provider: 'groq', apiKey: 'gsk-x' }, 'hello', 'SYS', { maxTokens: 1024 });

    const [url, init] = fetch.mock.calls[0];
    expect(url).toBe('https://api.groq.com/openai/v1/chat/completions');
    expect(init.headers).toEqual({ 'Content-Type': 'application/json', Authorization: 'Bearer gsk-x' });
    const body = JSON.parse(init.body);
    expect(body.model).toBe('llama-3.3-70b-versatile');
    expect(body.temperature).toBe(0.7);
    expect(body.max_tokens).toBe(1024);
  });

  it('OpenRouter: referer/title headers, NO temperature', async () => {
    const fetch = mockFetchOnce(openAiResponse);
    await callProvider({ provider: 'openrouter', apiKey: 'sk-or' }, 'hello', 'SYS', { maxTokens: 1024 });

    const [url, init] = fetch.mock.calls[0];
    expect(url).toBe('https://openrouter.ai/api/v1/chat/completions');
    expect(init.headers).toEqual({
      'Content-Type': 'application/json',
      Authorization: 'Bearer sk-or',
      'HTTP-Referer': 'https://lexai.dev',
      'X-Title': 'LexAI',
    });
    const body = JSON.parse(init.body);
    expect(body.model).toBe('openai/gpt-4o-mini');
    expect(body).not.toHaveProperty('temperature');
  });

  it('uses the configured model over the default', async () => {
    const fetch = mockFetchOnce(openAiResponse);
    await callProvider({ provider: 'openai', apiKey: 'k', model: 'gpt-4o' }, 'x', 'SYS');
    expect(JSON.parse(fetch.mock.calls[0][1].body).model).toBe('gpt-4o');
  });
});

describe('callProvider error handling (parity with old implementations)', () => {
  it('surfaces provider error messages with the provider label', async () => {
    // 401/403 additionally carry the settings hint — see 'key-rejection messaging'.
    mockFetchOnce({ error: { message: 'rate limited' } }, { ok: false, status: 429 });
    const res = await callProvider({ provider: 'openai', apiKey: 'bad' }, 'x', 'SYS');
    expect(res).toEqual({ error: 'OpenAI error: rate limited' });
  });

  it('falls back to HTTP status when the error body has no message', async () => {
    mockFetchOnce({}, { ok: false, status: 500 });
    const res = await callProvider({ provider: 'groq', apiKey: 'k' }, 'x', 'SYS');
    expect(res).toEqual({ error: 'Groq error: HTTP 500' });
  });

  it('reports network failures with the provider label', async () => {
    vi.stubGlobal('fetch', vi.fn().mockRejectedValue(new Error('offline')));
    const res = await callProvider({ provider: 'anthropic', apiKey: 'k' }, 'x', 'SYS');
    expect(res.error).toMatch(/^Network error reaching Anthropic:/);
  });

  it('reports empty responses', async () => {
    mockFetchOnce({ choices: [] });
    const res = await callProvider({ provider: 'openrouter', apiKey: 'k' }, 'x', 'SYS');
    expect(res).toEqual({ error: 'OpenRouter returned an empty response.' });
  });

  it('returns a clean HTTP error when the error body is not JSON (e.g. HTML 502)', async () => {
    vi.stubGlobal('fetch', vi.fn().mockResolvedValue({
      ok: false,
      status: 502,
      json: async () => { throw new SyntaxError('Unexpected token < in JSON'); },
    }));
    const res = await callProvider({ provider: 'openai', apiKey: 'k' }, 'x', 'SYS');
    expect(res).toEqual({ error: 'OpenAI error: HTTP 502' });
  });

  it('rejects unknown providers without fetching', async () => {
    const fetch = mockFetchOnce({});
    const res = await callProvider({ provider: 'bogus', apiKey: 'k' }, 'x', 'SYS');
    expect(res.error).toContain('Unknown provider: "bogus"');
    expect(fetch).not.toHaveBeenCalled();
  });

  it('defaults to openai when no provider is configured', async () => {
    const fetch = mockFetchOnce(openAiResponse);
    await callProvider({ apiKey: 'k' }, 'x', 'SYS');
    expect(fetch.mock.calls[0][0]).toBe('https://api.openai.com/v1/chat/completions');
  });
});

describe('getSystemPrompt', () => {
  it("normalizes 'fix' to the grammar prompt", () => {
    expect(getSystemPrompt('fix')).toBe(getSystemPrompt('grammar'));
    expect(getSystemPrompt('fix')).toContain('grammar editor');
  });

  it('falls back to grammar for unknown actions', () => {
    expect(getSystemPrompt('nonsense')).toBe(getSystemPrompt('grammar'));
  });

  it('appends a style modifier except for Default', () => {
    expect(getSystemPrompt('rephrase', 'Formal')).toMatch(/Write in a formal style\.$/);
    expect(getSystemPrompt('rephrase', 'Default')).not.toContain('style.');
    expect(getSystemPrompt('rephrase')).not.toContain('Write in a');
  });

  it("'prompt' uses the prompt-engineer prompt with a prompt-directed style modifier", () => {
    expect(getSystemPrompt('prompt')).toContain('expert prompt engineer');
    expect(getSystemPrompt('prompt', 'Formal')).toMatch(/instruct the model to respond in a formal style\.$/);
    expect(getSystemPrompt('prompt', 'Formal')).not.toContain('Write in a');
    expect(getSystemPrompt('prompt', 'Default')).toBe(getSystemPrompt('prompt'));
  });

  it('Prompt Builder params add instructions; Auto adds nothing', () => {
    const base = getSystemPrompt('prompt');
    expect(getSystemPrompt('prompt', undefined, { promptStyle: 'Auto', persona: 'Auto', format: 'Auto' })).toBe(base);

    const full = getSystemPrompt('prompt', undefined, {
      promptStyle: 'Few-shot',
      persona: 'Data Analyst',
      format: 'JSON',
    });
    expect(full).toContain('few-shot');
    expect(full).toContain('persona of Data Analyst');
    expect(full).toContain('final output as json');

    expect(getSystemPrompt('prompt', undefined, { persona: 'None' })).toContain('Do not assign a persona');
    // Params are prompt-action-only — other actions ignore them.
    expect(getSystemPrompt('rephrase', undefined, { persona: 'Teacher' })).toBe(getSystemPrompt('rephrase'));
  });
});

describe('defaultMaxTokens', () => {
  it('never goes below the old 1024 budget', () => {
    expect(defaultMaxTokens('short')).toBe(1024);
  });

  it('scales with input length and clamps at 8192', () => {
    expect(defaultMaxTokens('a'.repeat(4000))).toBe(4000);
    expect(defaultMaxTokens('a'.repeat(50000))).toBe(8192);
  });
});

describe('key-rejection messaging', () => {
  it('appends a settings hint to 401/403 chat errors only', async () => {
    mockFetchOnce({ error: { message: 'Invalid API Key' } }, { ok: false, status: 401 });
    const rejected = await callProvider({ provider: 'groq', apiKey: 'bad' }, 'hi', 'SYS');
    expect(rejected.error).toBe(
      'Groq error: Invalid API Key — open LexAI Settings and re-enter your API key for this provider.',
    );

    mockFetchOnce({ error: { message: 'server exploded' } }, { ok: false, status: 500 });
    const other = await callProvider({ provider: 'groq', apiKey: 'k' }, 'hi', 'SYS');
    expect(other.error).toBe('Groq error: server exploded');
  });

  it('maps provider ids to display names', () => {
    expect(providerLabel('groq')).toBe('Groq');
    expect(providerLabel('anthropic')).toBe('Anthropic');
    expect(providerLabel('mystery')).toBe('mystery');
  });
});

describe('listModels', () => {
  it('requires a key for Anthropic and sends the direct-browser-access header', async () => {
    expect(await listModels('anthropic')).toEqual({
      error: 'Anthropic requires an API key to list models.',
    });

    const fetch = mockFetchOnce({ data: [{ id: 'claude-3-5-haiku-20241022' }] });
    await listModels('anthropic', 'sk-ant');
    const [url, init] = fetch.mock.calls[0];
    expect(url).toBe('https://api.anthropic.com/v1/models');
    expect(init.headers['anthropic-dangerous-direct-browser-access']).toBe('true');
    expect(init.headers['x-api-key']).toBe('sk-ant');
  });

  it('allows keyless listing (OpenRouter) and sends bearer auth when a key exists', async () => {
    const noKey = mockFetchOnce({ data: [{ id: 'openai/gpt-4o' }] });
    await listModels('openrouter');
    expect(noKey.mock.calls[0][1].headers).not.toHaveProperty('Authorization');

    const withKey = mockFetchOnce({ data: [{ id: 'gpt-4o' }] });
    await listModels('openai', 'sk-x');
    expect(withKey.mock.calls[0][1].headers['Authorization']).toBe('Bearer sk-x');
  });

  it('filters non-chat models and sorts ids', async () => {
    mockFetchOnce({
      data: [
        { id: 'gpt-4o' },
        { id: 'text-embedding-3-small' },
        { id: 'whisper-1' },
        { id: 'dall-e-3' },
        { id: 'gpt-4o-mini' },
      ],
    });
    expect(await listModels('openai', 'k')).toEqual({ models: ['gpt-4o', 'gpt-4o-mini'] });
  });

  it('sends bearer auth for Groq and labels its errors with the display name', async () => {
    const fetch = mockFetchOnce({ data: [{ id: 'llama-3.3-70b-versatile' }] });
    expect(await listModels('groq', 'gsk_test')).toEqual({ models: ['llama-3.3-70b-versatile'] });
    const [url, init] = fetch.mock.calls[0];
    expect(url).toBe('https://api.groq.com/openai/v1/models');
    expect(init.headers['Authorization']).toBe('Bearer gsk_test');

    mockFetchOnce({ error: { message: 'Invalid API Key' } }, { ok: false, status: 401 });
    expect((await listModels('groq', 'bad')).error).toBe('Groq error: Invalid API Key');
  });

  it('flags a rejected key so Options can prompt for a new one', async () => {
    mockFetchOnce({ error: { message: 'Invalid API Key' } }, { ok: false, status: 401 });
    expect(await listModels('groq', 'bad')).toEqual({
      error: 'Groq error: Invalid API Key',
      keyRejected: true,
    });

    mockFetchOnce({ error: { message: 'boom' } }, { ok: false, status: 500 });
    expect((await listModels('groq', 'k')).keyRejected).toBeUndefined();
  });

  it('errors on empty lists and unknown providers', async () => {
    mockFetchOnce({ data: [] });
    expect((await listModels('openai', 'k')).error).toBe('No models returned by OpenAI.');
    expect((await listModels('bogus', 'k')).error).toContain('Unknown provider');
  });
});