feat(server): refactor for byok (#14911)

2026-08-23 13:02:21 +08:00 · 2026-05-07 04:03:14 +08:00
parent 4e169ea5c7
commit eb9cc22502
115 changed files with 10369 additions and 1256 deletions
@@ -0,0 +1,183 @@
+import path from 'node:path';
+
+import { app, safeStorage } from 'electron';
+
+import { PersistentJSONFileStorage } from '../shared-storage/json-file';
+import type { NamespaceHandlers } from '../type';
+
+const byokStorage = new PersistentJSONFileStorage(
+  path.join(app.getPath('userData'), 'workspace-byok-keys.json')
+);
+
+export function disposeWorkspaceByokStorage() {
+  byokStorage.dispose();
+}
+
+const allowedProviders = new Set(['openai', 'anthropic', 'gemini', 'fal']);
+
+type WorkspaceByokKey = {
+  id: string;
+  provider: 'openai' | 'anthropic' | 'gemini' | 'fal';
+  name: string;
+  description?: string | null;
+  apiKey: string;
+  endpoint?: string | null;
+  sortOrder?: number | null;
+  enabled?: boolean | null;
+};
+
+type WorkspaceByokKeyInput = Omit<WorkspaceByokKey, 'apiKey'> & {
+  apiKey?: string | null;
+};
+
+function assertSupported() {
+  if (!safeStorage.isEncryptionAvailable()) {
+    throw new Error('Secure BYOK key storage is not available.');
+  }
+}
+
+function hasOwnField(
+  key: WorkspaceByokKeyInput,
+  field: keyof WorkspaceByokKey
+) {
+  return Object.prototype.hasOwnProperty.call(key, field);
+}
+
+function normalizeKey(
+  key: WorkspaceByokKeyInput,
+  existing?: WorkspaceByokKey,
+  defaultSortOrder = 0
+): WorkspaceByokKey {
+  if (!allowedProviders.has(key.provider)) {
+    throw new Error('Unsupported BYOK provider.');
+  }
+  const apiKey = key.apiKey ?? existing?.apiKey;
+  if (!key.id || !key.name || !apiKey) {
+    throw new Error('Invalid BYOK key.');
+  }
+  return {
+    id: key.id,
+    provider: key.provider,
+    name: key.name,
+    description: hasOwnField(key, 'description')
+      ? (key.description ?? null)
+      : (existing?.description ?? null),
+    apiKey,
+    endpoint: hasOwnField(key, 'endpoint')
+      ? (key.endpoint ?? null)
+      : (existing?.endpoint ?? null),
+    sortOrder: hasOwnField(key, 'sortOrder')
+      ? (key.sortOrder ?? defaultSortOrder)
+      : (existing?.sortOrder ?? defaultSortOrder),
+    enabled: hasOwnField(key, 'enabled')
+      ? (key.enabled ?? true)
+      : (existing?.enabled ?? true),
+  };
+}
+
+function encryptKey(key: WorkspaceByokKey) {
+  return safeStorage
+    .encryptString(JSON.stringify(normalizeKey(key)))
+    .toString('base64');
+}
+
+function decryptKey(value: string): WorkspaceByokKey | null {
+  try {
+    return normalizeKey(
+      JSON.parse(safeStorage.decryptString(Buffer.from(value, 'base64')))
+    );
+  } catch {
+    return null;
+  }
+}
+
+function sortWorkspaceKeys(keys: WorkspaceByokKey[]) {
+  return keys.toSorted((a, b) => (a.sortOrder ?? 0) - (b.sortOrder ?? 0));
+}
+
+function readWorkspaceKeys(workspaceId: string): WorkspaceByokKey[] {
+  assertSupported();
+  const encryptedKeys = byokStorage.get<string[]>(workspaceId) ?? [];
+  return sortWorkspaceKeys(
+    encryptedKeys.flatMap(value => {
+      const key = decryptKey(value);
+      return key ? [key] : [];
+    })
+  );
+}
+
+function writeWorkspaceKeys(workspaceId: string, keys: WorkspaceByokKey[]) {
+  assertSupported();
+  byokStorage.set(workspaceId, keys.map(encryptKey));
+}
+
+function toPublicKey({ apiKey: _, ...key }: WorkspaceByokKey) {
+  return {
+    ...key,
+    storage: 'local',
+    configured: true,
+    endpointEditable: false,
+    testStatus: 'passed',
+  };
+}
+
+export const byokStorageHandlers = {
+  isSupported: async () => safeStorage.isEncryptionAvailable(),
+  listWorkspaceKeys: async (_e, workspaceId: string) => {
+    return readWorkspaceKeys(workspaceId).map(toPublicKey);
+  },
+  getWorkspaceLeaseProviders: async (_e, workspaceId: string) => {
+    return readWorkspaceKeys(workspaceId).filter(key => key.enabled !== false);
+  },
+  upsertWorkspaceKey: async (
+    _e,
+    workspaceId: string,
+    key: WorkspaceByokKeyInput
+  ) => {
+    const keys = readWorkspaceKeys(workspaceId);
+    const index = keys.findIndex(storedKey => storedKey.id === key.id);
+    const nextKey = normalizeKey(
+      key,
+      index === -1 ? undefined : keys[index],
+      keys.length
+    );
+    if (index === -1) {
+      keys.push(nextKey);
+    } else {
+      keys[index] = nextKey;
+    }
+    writeWorkspaceKeys(workspaceId, keys);
+    return toPublicKey(nextKey);
+  },
+  deleteWorkspaceKey: async (_e, workspaceId: string, keyId: string) => {
+    writeWorkspaceKeys(
+      workspaceId,
+      readWorkspaceKeys(workspaceId).filter(key => key.id !== keyId)
+    );
+    return true;
+  },
+  reorderWorkspaceKeys: async (_e, workspaceId: string, ids: string[]) => {
+    const keys = readWorkspaceKeys(workspaceId);
+    const byId = new Map(keys.map(key => [key.id, key]));
+    const ordered = ids
+      .map((id, sortOrder) => {
+        const key = byId.get(id);
+        byId.delete(id);
+        return key ? ({ ...key, sortOrder } as WorkspaceByokKey) : null;
+      })
+      .filter((key): key is WorkspaceByokKey => !!key);
+    const nextKeys = sortWorkspaceKeys([
+      ...ordered,
+      ...Array.from(byId.values()).map((key, index) => ({
+        ...key,
+        sortOrder: ordered.length + index,
+      })),
+    ]);
+    writeWorkspaceKeys(workspaceId, nextKeys);
+    return nextKeys.map(toPublicKey);
+  },
+  clearWorkspaceKeys: async (_e, workspaceId: string) => {
+    byokStorage.del(workspaceId);
+    return true;
+  },
+} satisfies NamespaceHandlers;
@@ -2,6 +2,7 @@ import { I18n } from '@affine/i18n';
 import { ipcMain } from 'electron';

 import { AFFINE_API_CHANNEL_NAME } from '../shared/type';
+import { byokStorageHandlers } from './byok-storage/handlers';
 import { clipboardHandlers } from './clipboard';
 import { configStorageHandlers } from './config-storage';
 import { findInPageHandlers } from './find-in-page';
@@ -42,6 +43,7 @@ export const allHandlers = {
  recording: recordingHandlers,
  popup: popupHandlers,
  i18n: i18nHandlers,
+  byokStorage: byokStorageHandlers,
 };

 export const registerHandlers = () => {
@@ -0,0 +1,151 @@
+import path from 'node:path';
+
+import fs from 'fs-extra';
+import { afterEach, beforeEach, describe, expect, test, vi } from 'vitest';
+
+const tmpDir = path.join(__dirname, 'tmp-byok-storage');
+let disposeWorkspaceByokStorage: (() => void) | undefined;
+
+vi.mock('electron', () => ({
+  app: {
+    getPath: () => tmpDir,
+    on: vi.fn(),
+  },
+  safeStorage: {
+    isEncryptionAvailable: () => true,
+    encryptString: (value: string) => Buffer.from(value, 'utf-8'),
+    decryptString: (value: Buffer) => value.toString('utf-8'),
+  },
+}));
+
+beforeEach(async () => {
+  vi.resetModules();
+  disposeWorkspaceByokStorage = undefined;
+  await fs.remove(tmpDir);
+});
+
+afterEach(async () => {
+  disposeWorkspaceByokStorage?.();
+  vi.resetModules();
+  await fs.remove(tmpDir);
+});
+
+describe('byok storage handlers', () => {
+  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',
+      apiKey: 'sk-openai',
+      sortOrder: 1,
+    });
+    await byokStorageHandlers.upsertWorkspaceKey(ipcEvent, 'workspace-1', {
+      id: 'local-gemini',
+      provider: 'gemini',
+      name: 'Gemini',
+      apiKey: 'sk-gemini',
+      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.apiKey)).toEqual([
+      'sk-openai',
+      'sk-gemini',
+    ]);
+
+    await byokStorageHandlers.clearWorkspaceKeys(ipcEvent, 'workspace-1');
+    await expect(
+      byokStorageHandlers.listWorkspaceKeys(ipcEvent, 'workspace-1')
+    ).resolves.toEqual([]);
+  });
+
+  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',
+      apiKey: 'sk-openai',
+      endpoint: 'https://api.openai.example/v1',
+      sortOrder: 4,
+      enabled: false,
+    });
+
+    await byokStorageHandlers.upsertWorkspaceKey(ipcEvent, 'workspace-1', {
+      id: 'local-openai',
+      provider: 'openai',
+      name: 'OpenAI renamed',
+      apiKey: 'sk-openai-next',
+    });
+
+    const [publicKey] = await byokStorageHandlers.listWorkspaceKeys(
+      ipcEvent,
+      'workspace-1'
+    );
+    expect(publicKey).toMatchObject({
+      id: 'local-openai',
+      name: 'OpenAI renamed',
+      description: 'Primary key',
+      endpoint: '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',
+      apiKey: 'sk-openai-next',
+      endpoint: 'https://api.openai.example/v1',
+      sortOrder: 4,
+      enabled: true,
+    });
+  });
+});
@@ -23,7 +23,8 @@ export default defineConfig({
  test: {
    setupFiles: [resolve(rootDir, './scripts/setup/global.ts')],
    include: ['./test/**/*.spec.ts'],
-    testTimeout: 30000,
+    testTimeout: 60000,
+    hookTimeout: 30000,
    pool: 'forks',
    maxWorkers: 1,
    coverage: {
@@ -1,57 +0,0 @@
-// @generated
-// This file was automatically generated and should not be edited.
-
-@_exported import ApolloAPI
-
-public class ApplyDocUpdatesMutation: GraphQLMutation {
-  public static let operationName: String = "applyDocUpdates"
-  public static let operationDocument: ApolloAPI.OperationDocument = .init(
-    definition: .init(
-      #"mutation applyDocUpdates($workspaceId: String!, $docId: String!, $op: String!, $updates: String!) { applyDocUpdates( workspaceId: $workspaceId docId: $docId op: $op updates: $updates ) }"#
-    ))
-
-  public var workspaceId: String
-  public var docId: String
-  public var op: String
-  public var updates: String
-
-  public init(
-    workspaceId: String,
-    docId: String,
-    op: String,
-    updates: String
-  ) {
-    self.workspaceId = workspaceId
-    self.docId = docId
-    self.op = op
-    self.updates = updates
-  }
-
-  public var __variables: Variables? { [
-    "workspaceId": workspaceId,
-    "docId": docId,
-    "op": op,
-    "updates": updates
-  ] }
-
-  public struct Data: AffineGraphQL.SelectionSet {
-    public let __data: DataDict
-    public init(_dataDict: DataDict) { __data = _dataDict }
-
-    public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.Mutation }
-    public static var __selections: [ApolloAPI.Selection] { [
-      .field("applyDocUpdates", String.self, arguments: [
-        "workspaceId": .variable("workspaceId"),
-        "docId": .variable("docId"),
-        "op": .variable("op"),
-        "updates": .variable("updates")
-      ]),
-    ] }
-    public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [
-      ApplyDocUpdatesMutation.Data.self
-    ] }
-
-    /// Apply updates to a doc using LLM and return the merged markdown.
-    public var applyDocUpdates: String { __data["applyDocUpdates"] }
-  }
-}