makelore/tests/unit/pi-provider-config.test.ts

import { mkdtemp, readFile, rm, writeFile } from 'node:fs/promises';
import { tmpdir } from 'node:os';
import path from 'node:path';
import { afterEach, describe, expect, it, vi } from 'vitest';
import type { ProviderAccount } from '@electron/shared/providers/types';
import {
  buildPiProviderCatalog,
  buildPiWorkerCredentialProjection,
  credentialValueForProviderSecret,
  PiProviderConfigError,
  resolvePiRuntimeProviderId,
  selectPiProviderModel,
  summarizePiWorkerCredentialProjection,
  writePiProviderCatalog,
} from '@electron/coding-runtime/pi/provider-config';

const temporaryRoots: string[] = [];

afterEach(async () => {
  await Promise.all(temporaryRoots.splice(0).map((root) => rm(root, { recursive: true, force: true })));
});

function account(overrides: Partial<ProviderAccount> = {}): ProviderAccount {
  return {
    id: 'account-one',
    vendorId: 'openai',
    label: 'Primary account',
    authMode: 'api_key',
    model: 'gpt-5.4',
    enabled: true,
    isDefault: true,
    createdAt: '2026-08-22T00:00:00.000Z',
    updatedAt: '2026-08-22T00:00:00.000Z',
    ...overrides,
  };
}

describe('Pi Provider catalog', () => {
  it('maps every current Provider protocol shape to a Pi API', () => {
    const catalog = buildPiProviderCatalog({
      accounts: [
        account({ id: 'completions', vendorId: 'custom', apiProtocol: 'openai-completions', baseUrl: 'https://one.test/v1', model: 'chat' }),
        account({ id: 'responses', vendorId: 'custom', apiProtocol: 'openai-responses', baseUrl: 'https://two.test/v1', model: 'response' }),
        account({ id: 'anthropic', vendorId: 'anthropic', model: 'claude-opus-4-6' }),
        account({ id: 'google', vendorId: 'google', model: 'gemini-3-pro-preview' }),
        account({ id: 'openrouter', vendorId: 'openrouter', model: 'openai/gpt-5.4' }),
      ],
    });

    expect(catalog.descriptors.map(({ api }) => api)).toEqual([
      'openai-completions',
      'openai-responses',
      'anthropic-messages',
      'google-generative-ai',
      'openai-completions',
    ]);
    expect(catalog.descriptors.at(-1)?.models[0]?.compat).toEqual({
      thinkingFormat: 'openrouter',
      sessionAffinityFormat: 'openrouter',
    });
    expect(Object.keys(catalog.descriptors.at(-1)?.headers ?? {}).sort()).toEqual([
      'HTTP-Referer',
      'X-OpenRouter-Title',
    ]);
  });

  it('uses stable collision-free account-derived runtime IDs for the same vendor', () => {
    const first = account({ id: 'openai-account-a' });
    const second = account({ id: 'openai-account-b', isDefault: false });
    const catalog = buildPiProviderCatalog({ accounts: [first, second] });

    expect(catalog.descriptors[0]?.runtimeProviderId).toBe(resolvePiRuntimeProviderId(first.id));
    expect(catalog.descriptors[1]?.runtimeProviderId).toBe(resolvePiRuntimeProviderId(second.id));
    expect(catalog.descriptors[0]?.runtimeProviderId).not.toBe(catalog.descriptors[1]?.runtimeProviderId);
  });

  it('normalizes Works gateway /v1 and keeps all credential and header values out of catalog output', async () => {
    const provider = account({
      id: 'niancode-user-models',
      vendorId: 'custom',
      apiProtocol: 'openai-completions',
      baseUrl: 'https://gateway.test/',
      model: 'qwen3.6-plus',
      headers: {
        'X-Works-Square-AI-Token': '{env:OLD_GATEWAY_TOKEN}',
        'X-Tenant-Secret': 'private-tenant-header',
      },
      metadata: {
        worksSquareCredentialMode: 'works_square_ai_gateway',
        customModels: ['qwen3.6-plus'],
      },
    });
    const catalog = buildPiProviderCatalog({ accounts: [provider] });
    const descriptor = catalog.descriptors[0]!;
    const serialized = JSON.stringify(catalog);

    expect(descriptor.baseUrl).toBe('https://gateway.test/v1');
    expect(Object.keys(descriptor.headers).sort()).toEqual([
      'Authorization',
      'X-Tenant-Secret',
      'X-Works-Square-AI-Token',
    ]);
    expect(descriptor.models[0]).toMatchObject({
      id: 'qwen3.6-plus',
      input: ['text', 'image'],
      contextWindow: 1_000_000,
      maxOutputTokens: 65_536,
    });
    expect(serialized).not.toContain('private-tenant-header');
    expect(serialized).not.toContain('OLD_GATEWAY_TOKEN');

    const projection = await buildPiWorkerCredentialProjection({
      account: provider,
      descriptor,
      resolveCredential: vi.fn().mockResolvedValue('gateway-proxy-token'),
    });
    expect(Object.values(projection.env)).toEqual(expect.arrayContaining([
      'gateway-proxy-token',
      'Bearer gateway-proxy-token',
      'private-tenant-header',
    ]));
    expect(projection.sensitiveValues).toContain('gateway-proxy-token');
    const safeProjection = summarizePiWorkerCredentialProjection(projection);
    expect(JSON.stringify(safeProjection)).not.toContain('gateway-proxy-token');
    expect(JSON.stringify(safeProjection)).not.toContain('private-tenant-header');
  });

  it('uses only the current worker local-proxy credential for proxy mode', async () => {
    const provider = account({
      id: 'niancode-user-models',
      vendorId: 'custom',
      apiProtocol: 'openai-completions',
      baseUrl: 'http://127.0.0.1:54321/api/ai-proxy/v1',
      model: 'qwen-vl-max',
      metadata: { worksSquareCredentialMode: 'works_square_ai_gateway_proxy' },
    });
    const descriptor = buildPiProviderCatalog({ accounts: [provider] }).descriptors[0]!;
    const resolveCredential = vi.fn().mockResolvedValue('stale-stored-host-token');
    const projection = await buildPiWorkerCredentialProjection({
      account: provider,
      descriptor,
      resolveCredential,
      localProxyCredential: 'current-worker-host-token',
    });

    expect(resolveCredential).not.toHaveBeenCalled();
    expect(Object.values(projection.env)).toContain('current-worker-host-token');
    expect(Object.values(projection.env)).not.toContain('stale-stored-host-token');
  });

  it('uses account-scoped model capability metadata and rejects unavailable models', () => {
    const provider = account({ id: 'account-with-vision', model: 'vision-model' });
    const catalog = buildPiProviderCatalog({
      accounts: [provider],
      modelSummaries: [{
        id: 'vision-model',
        name: 'Vision model',
        vendorId: 'openai',
        accountId: provider.id,
        supportsVision: true,
        supportsReasoning: true,
        contextWindow: 200_000,
        source: 'remote',
      }],
    });
    expect(selectPiProviderModel(catalog, {
      accountId: provider.id,
      modelId: 'vision-model',
      thinkingLevel: 'high',
    })).toMatchObject({
      input: ['text', 'image'],
      contextWindow: 200_000,
      thinkingLevel: 'high',
    });
    expect(() => selectPiProviderModel(catalog, {
      accountId: provider.id,
      modelId: 'missing-model',
      thinkingLevel: 'off',
    })).toThrowError(PiProviderConfigError);
  });

  it('does not replace an existing catalog when model selection is unavailable', async () => {
    const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-provider-'));
    temporaryRoots.push(root);
    const filePath = path.join(root, 'models.json');
    await writeFile(filePath, 'existing-catalog\n', 'utf8');
    const catalog = buildPiProviderCatalog({ accounts: [account()] });

    await expect(writePiProviderCatalog(filePath, catalog, {
      accountId: 'account-one',
      modelId: 'not-configured',
      thinkingLevel: 'off',
    })).rejects.toMatchObject({ code: 'MODEL_UNAVAILABLE' });
    expect(await readFile(filePath, 'utf8')).toBe('existing-catalog\n');
  });

  it('fails closed when a non-local credential is unavailable', async () => {
    const provider = account();
    const descriptor = buildPiProviderCatalog({ accounts: [provider] }).descriptors[0]!;
    await expect(buildPiWorkerCredentialProjection({
      account: provider,
      descriptor,
      resolveCredential: vi.fn().mockResolvedValue(null),
    })).rejects.toMatchObject({ code: 'PROVIDER_AUTH_REQUIRED' });
  });

  it('projects API key, OAuth, and local secrets to one worker credential value', () => {
    expect(credentialValueForProviderSecret({
      type: 'api_key',
      accountId: 'account-one',
      apiKey: 'api-key-value',
    })).toBe('api-key-value');
    expect(credentialValueForProviderSecret({
      type: 'oauth',
      accountId: 'account-one',
      accessToken: 'oauth-access-token',
      refreshToken: 'refresh-token-never-projected',
      expiresAt: Date.now() + 60_000,
    })).toBe('oauth-access-token');
    expect(credentialValueForProviderSecret({
      type: 'local',
      accountId: 'account-one',
    })).toBeNull();
  });
});