AFFiNE-Mirror/packages/frontend/apps/electron/test/main/byok-storage.spec.ts

import os from 'node:os';
import path from 'node:path';

import fs from 'fs-extra';
import {
  afterEach,
  beforeAll,
  beforeEach,
  describe,
  expect,
  test,
  vi,
} from 'vitest';

const electronMock = vi.hoisted(() => ({
  tmpDir: '',
  appOn: vi.fn(),
  isEncryptionAvailable: vi.fn(() => true),
  encryptString: vi.fn((value: string) => Buffer.from(value, 'utf-8')),
  decryptString: vi.fn((value: Buffer) => value.toString('utf-8')),
}));

let disposeWorkspaceByokStorage: (() => void) | undefined;

vi.mock('electron', () => ({
  app: {
    getPath: () => electronMock.tmpDir,
    on: electronMock.appOn,
  },
  safeStorage: {
    isEncryptionAvailable: electronMock.isEncryptionAvailable,
    encryptString: electronMock.encryptString,
    decryptString: electronMock.decryptString,
  },
}));

vi.mock('../../src/main/logger', () => ({
  logger: {
    error: vi.fn(),
  },
}));

// Warm the handler module's transform once so the first per-test dynamic
// import doesn't race the default 60s test timeout on loaded CI shards, where
// cold-transforming the heavy `@toeverything/infra` graph can starve.
beforeAll(async () => {
  await import('@affine/electron/main/byok-storage/handlers');
}, 120_000);

beforeEach(async () => {
  vi.useRealTimers();
  vi.resetModules();
  electronMock.appOn.mockReset();
  electronMock.isEncryptionAvailable.mockReset().mockReturnValue(true);
  electronMock.encryptString
    .mockReset()
    .mockImplementation((value: string) => Buffer.from(value, 'utf-8'));
  electronMock.decryptString
    .mockReset()
    .mockImplementation((value: Buffer) => value.toString('utf-8'));
  disposeWorkspaceByokStorage = undefined;
  electronMock.tmpDir = await fs.mkdtemp(
    path.join(os.tmpdir(), 'affine-byok-storage-')
  );
});

afterEach(async () => {
  disposeWorkspaceByokStorage?.();
  disposeWorkspaceByokStorage = undefined;
  vi.resetModules();
  if (electronMock.tmpDir) {
    await fs.remove(electronMock.tmpDir);
  }
  electronMock.tmpDir = '';
});

describe('byok storage handlers', () => {
  const definition = {
    version: 1,
    endpoint: { kind: 'provider_default' },
    models: [
      {
        modelId: 'model-1',
        enabled: true,
        capabilities: [
          {
            input: ['text'],
            output: ['text'],
            features: [],
            attachmentKinds: [],
            attachmentSources: [],
          },
        ],
      },
    ],
  };
  test('stores encrypted local keys and keeps lease providers sorted', async () => {
    const { byokStorageHandlers, disposeWorkspaceByokStorage: dispose } =
      await import('@affine/electron/main/byok-storage/handlers');
    disposeWorkspaceByokStorage = dispose;
    const ipcEvent = undefined;

    await byokStorageHandlers.upsertWorkspaceKey(ipcEvent, 'workspace-1', {
      id: 'local-openai',
      provider: 'openai',
      name: 'OpenAI',
      credential: 'sk-openai',
      definition,
      sortOrder: 1,
    });
    await byokStorageHandlers.upsertWorkspaceKey(ipcEvent, 'workspace-1', {
      id: 'local-gemini',
      provider: 'gemini',
      name: 'Gemini',
      credential: 'sk-gemini',
      definition,
      sortOrder: 0,
    });

    const list = await byokStorageHandlers.listWorkspaceKeys(
      ipcEvent,
      'workspace-1'
    );
    expect(list.map(key => key.id)).toEqual(['local-gemini', 'local-openai']);
    expect(JSON.stringify(list)).not.toContain('sk-openai');

    const reordered = await byokStorageHandlers.reorderWorkspaceKeys(
      ipcEvent,
      'workspace-1',
      ['local-openai', 'local-gemini']
    );
    expect(reordered.map(key => key.id)).toEqual([
      'local-openai',
      'local-gemini',
    ]);

    const leaseProviders = await byokStorageHandlers.getWorkspaceLeaseProviders(
      ipcEvent,
      'workspace-1'
    );
    expect(leaseProviders.map(key => key.credential)).toEqual([
      'sk-openai',
      'sk-gemini',
    ]);

    await byokStorageHandlers.clearWorkspaceKeys(ipcEvent, 'workspace-1');
    await expect(
      byokStorageHandlers.listWorkspaceKeys(ipcEvent, 'workspace-1')
    ).resolves.toEqual([]);
  });

  test('does not write local keys when secure storage is unavailable', async () => {
    electronMock.isEncryptionAvailable.mockReturnValue(false);

    const { byokStorageHandlers, disposeWorkspaceByokStorage: dispose } =
      await import('@affine/electron/main/byok-storage/handlers');
    disposeWorkspaceByokStorage = dispose;
    const ipcEvent = undefined;

    await expect(byokStorageHandlers.isSupported()).resolves.toBe(false);
    await expect(
      byokStorageHandlers.upsertWorkspaceKey(ipcEvent, 'workspace-1', {
        id: 'local-openai',
        provider: 'openai',
        name: 'OpenAI',
        credential: 'sk-openai',
        definition,
      })
    ).rejects.toThrow('Secure BYOK key storage is not available.');
    expect(electronMock.encryptString).not.toHaveBeenCalled();
  });

  test.each([
    [
      'custom endpoint without URL',
      { ...definition, endpoint: { kind: 'custom' } },
    ],
    [
      'unsupported endpoint protocol',
      { ...definition, endpoint: { kind: 'custom', url: 'file:///tmp/api' } },
    ],
    [
      'malformed capability object',
      {
        ...definition,
        models: [{ ...definition.models[0], capabilities: [{}] }],
      },
    ],
    [
      'unknown capability value',
      {
        ...definition,
        models: [
          {
            ...definition.models[0],
            capabilities: [
              { ...definition.models[0].capabilities[0], input: ['video'] },
            ],
          },
        ],
      },
    ],
  ])('rejects %s from IPC input', async (_name, malformedDefinition) => {
    const { byokStorageHandlers, disposeWorkspaceByokStorage: dispose } =
      await import('@affine/electron/main/byok-storage/handlers');
    disposeWorkspaceByokStorage = dispose;

    await expect(
      byokStorageHandlers.upsertWorkspaceKey(undefined, 'workspace-1', {
        id: 'local-openai',
        provider: 'openai',
        name: 'OpenAI',
        credential: 'sk-openai',
        definition: malformedDefinition as typeof definition,
      })
    ).rejects.toThrow('Invalid BYOK key.');
    expect(electronMock.encryptString).not.toHaveBeenCalled();
  });

  test('preserves existing local key fields during partial updates', async () => {
    const { byokStorageHandlers, disposeWorkspaceByokStorage: dispose } =
      await import('@affine/electron/main/byok-storage/handlers');
    disposeWorkspaceByokStorage = dispose;
    const ipcEvent = undefined;

    await byokStorageHandlers.upsertWorkspaceKey(ipcEvent, 'workspace-1', {
      id: 'local-openai',
      provider: 'openai',
      name: 'OpenAI',
      description: 'Primary key',
      credential: 'sk-openai',
      definition: {
        ...definition,
        endpoint: { kind: 'custom', url: 'https://api.openai.example/v1' },
      },
      sortOrder: 4,
      enabled: false,
    });

    await byokStorageHandlers.upsertWorkspaceKey(ipcEvent, 'workspace-1', {
      id: 'local-openai',
      provider: 'openai',
      name: 'OpenAI renamed',
      credential: 'sk-openai-next',
    });

    const [publicKey] = await byokStorageHandlers.listWorkspaceKeys(
      ipcEvent,
      'workspace-1'
    );
    expect(publicKey).toMatchObject({
      id: 'local-openai',
      name: 'OpenAI renamed',
      description: 'Primary key',
      definition: {
        ...definition,
        endpoint: { kind: 'custom', url: 'https://api.openai.example/v1' },
      },
      sortOrder: 4,
      enabled: false,
    });

    const [leaseProvider] =
      await byokStorageHandlers.getWorkspaceLeaseProviders(
        ipcEvent,
        'workspace-1'
      );
    expect(leaseProvider).toBeUndefined();

    await byokStorageHandlers.upsertWorkspaceKey(ipcEvent, 'workspace-1', {
      id: 'local-openai',
      provider: 'openai',
      name: 'OpenAI renamed again',
      enabled: true,
    });

    const [enabledLeaseProvider] =
      await byokStorageHandlers.getWorkspaceLeaseProviders(
        ipcEvent,
        'workspace-1'
      );
    expect(enabledLeaseProvider).toMatchObject({
      name: 'OpenAI renamed again',
      credential: 'sk-openai-next',
      definition: {
        ...definition,
        endpoint: { kind: 'custom', url: 'https://api.openai.example/v1' },
      },
      sortOrder: 4,
      enabled: true,
    });
  });
});