snapshot of https://github.com/toeverything/AFFiNE.git (canary, 2026-09-16; blobs > 131072B omitted)

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

@@ -0,0 +1,314 @@
+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 = {
+    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: 'openai_compatible', dialect: 'responses' },
+      },
+    ],
+    [
+      'unsupported endpoint protocol',
+      {
+        ...definition,
+        endpoint: {
+          kind: 'openai_compatible',
+          url: 'file:///tmp/api',
+          dialect: 'responses',
+        },
+      },
+    ],
+    [
+      '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: 'openai_compatible',
+          url: 'https://api.openai.example/v1',
+          dialect: 'responses',
+        },
+      },
+      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: 'openai_compatible',
+          url: 'https://api.openai.example/v1',
+          dialect: 'responses',
+        },
+      },
+      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: 'openai_compatible',
+          url: 'https://api.openai.example/v1',
+          dialect: 'responses',
+        },
+      },
+      sortOrder: 4,
+      enabled: true,
+    });
+  });
+});