@lobehub/chat 1.36.46 → 1.37.1

package/src/services/topic/pglite.test.ts

@@ -0,0 +1,212 @@
+import { eq } from 'drizzle-orm';
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { clientDB, initializeDB } from '@/database/client/db';
+import { sessions, topics, users } from '@/database/schemas';
+import { ChatTopic } from '@/types/topic';
+
+import { ClientService } from './pglite';
+
+// Mock data
+const userId = 'topic-user-test';
+const sessionId = 'topic-session';
+const mockTopicId = 'mock-topic-id';
+
+const mockTopic = {
+  id: mockTopicId,
+  title: 'Mock Topic',
+};
+
+const topicService = new ClientService(userId);
+
+beforeEach(async () => {
+  await initializeDB();
+
+  await clientDB.delete(users);
+
+  // 创建测试数据
+  await clientDB.transaction(async (tx) => {
+    await tx.insert(users).values({ id: userId });
+    await tx.insert(sessions).values({ id: sessionId, userId });
+    await tx.insert(topics).values({ ...mockTopic, sessionId, userId });
+  });
+});
+
+describe('TopicService', () => {
+  describe('createTopic', () => {
+    it('should create a topic and return its id', async () => {
+      // Setup
+      const createParams = {
+        title: 'New Topic',
+        sessionId: sessionId,
+      };
+
+      // Execute
+      const topicId = await topicService.createTopic(createParams);
+
+      // Assert
+      expect(topicId).toBeDefined();
+    });
+
+    it('should throw an error if topic creation fails', async () => {
+      // Setup
+      const createParams = {
+        title: 'New Topic',
+        sessionId: 123 as any, // sessionId should be string
+      };
+
+      // Execute & Assert
+      await expect(topicService.createTopic(createParams)).rejects.toThrowError();
+    });
+  });
+
+  describe('getTopics', () => {
+    // Example for getTopics
+    it('should query topics with given parameters', async () => {
+      // Setup
+      const queryParams = { sessionId };
+
+      // Execute
+      const data = await topicService.getTopics(queryParams);
+
+      // Assert
+      expect(data[0]).toMatchObject(mockTopic);
+    });
+  });
+
+  describe('updateTopic', () => {
+    // Example for updateFavorite
+    it('should toggle favorite status of a topic', async () => {
+      // Execute
+      const result = await topicService.updateTopic(mockTopicId, { favorite: true });
+
+      // Assert
+      expect(result[0].favorite).toBeTruthy();
+    });
+
+    it('should update the title of a topic', async () => {
+      // Setup
+      const newTitle = 'Updated Topic Title';
+
+      // Execute
+      const result = await topicService.updateTopic(mockTopicId, { title: newTitle });
+
+      // Assert
+      expect(result[0].title).toEqual(newTitle);
+    });
+  });
+
+  describe('removeTopic', () => {
+    it('should remove a topic by id', async () => {
+      // Execute
+      await topicService.removeTopic(mockTopicId);
+      const result = await clientDB.query.topics.findFirst({ where: eq(topics.id, mockTopicId) });
+
+      // Assert
+      expect(result).toBeUndefined();
+    });
+  });
+
+  describe('removeTopics', () => {
+    it('should remove all topics with a given session id', async () => {
+      // Setup
+      const sessionId = 'session-id';
+
+      // Execute
+      await topicService.removeTopics(sessionId);
+      const result = await clientDB.query.topics.findMany({
+        where: eq(topics.sessionId, sessionId),
+      });
+
+      expect(result.length).toEqual(0);
+    });
+  });
+
+  describe('batchRemoveTopics', () => {
+    it('should batch remove topics', async () => {
+      await clientDB.insert(topics).values([{ id: 'topic-id-1', title: 'topic-title', userId }]);
+      // Setup
+      const topicIds = [mockTopicId, 'another-topic-id'];
+
+      // Execute
+      await topicService.batchRemoveTopics(topicIds);
+
+      const count = await clientDB.$count(topics);
+
+      // Assert
+      expect(count).toBe(1);
+    });
+  });
+
+  describe('removeAllTopic', () => {
+    it('should clear all topics from the table', async () => {
+      // Execute
+      await topicService.removeAllTopic();
+
+      const count = await clientDB.$count(topics);
+      // Assert
+      expect(count).toBe(0);
+    });
+  });
+
+  describe('batchCreateTopics', () => {
+    it('should batch create topics', async () => {
+      // Execute
+      const result = await topicService.batchCreateTopics([
+        { id: 'topic-id-1', title: 'topic-title' },
+        { id: 'topic-id-2', title: 'topic-title' },
+      ] as ChatTopic[]);
+
+      // Assert
+      expect(result.success).toBeTruthy();
+      expect(result.added).toBe(2);
+    });
+  });
+
+  describe('getAllTopics', () => {
+    it('should retrieve all topics', async () => {
+      await clientDB.insert(topics).values([
+        { id: 'topic-id-1', title: 'topic-title', userId },
+        { id: 'topic-id-2', title: 'topic-title', userId },
+      ]);
+      // Execute
+      const result = await topicService.getAllTopics();
+
+      // Assert
+      expect(result.length).toEqual(3);
+    });
+  });
+
+  describe('searchTopics', () => {
+    it('should return all topics that match the keyword', async () => {
+      // Setup
+      const keyword = 'Topic';
+
+      // Execute
+      const result = await topicService.searchTopics(keyword, sessionId);
+
+      // Assert
+      expect(result.length).toEqual(1);
+    });
+    it('should return empty topic if not match the keyword', async () => {
+      // Setup
+      const keyword = 'search';
+
+      // Execute
+      const result = await topicService.searchTopics(keyword, sessionId);
+
+      // Assert
+      expect(result.length).toEqual(0);
+    });
+  });
+
+  describe('countTopics', () => {
+    it('should return topic counts', async () => {
+      // Execute
+      const result = await topicService.countTopics();
+
+      // Assert
+      expect(result).toBe(1);
+    });
+  });
+});

package/src/services/topic/pglite.ts

@@ -0,0 +1,85 @@
+import { INBOX_SESSION_ID } from '@/const/session';
+import { clientDB } from '@/database/client/db';
+import { TopicModel } from '@/database/server/models/topic';
+import { BaseClientService } from '@/services/baseClientService';
+import { ChatTopic } from '@/types/topic';
+
+import { CreateTopicParams, ITopicService, QueryTopicParams } from './type';
+
+export class ClientService extends BaseClientService implements ITopicService {
+  private get topicModel(): TopicModel {
+    return new TopicModel(clientDB as any, this.userId);
+  }
+
+  async createTopic(params: CreateTopicParams): Promise<string> {
+    const item = await this.topicModel.create({
+      ...params,
+      sessionId: this.toDbSessionId(params.sessionId),
+    } as any);
+
+    if (!item) {
+      throw new Error('topic create Error');
+    }
+
+    return item.id;
+  }
+
+  async batchCreateTopics(importTopics: ChatTopic[]) {
+    const data = await this.topicModel.batchCreate(importTopics as any);
+
+    return { added: data.length, ids: [], skips: [], success: true };
+  }
+
+  async cloneTopic(id: string, newTitle?: string) {
+    const data = await this.topicModel.duplicate(id, newTitle);
+    return data.topic.id;
+  }
+
+  async getTopics(params: QueryTopicParams) {
+    const data = await this.topicModel.query({
+      ...params,
+      sessionId: this.toDbSessionId(params.sessionId),
+    });
+    return data as unknown as Promise<ChatTopic[]>;
+  }
+
+  async searchTopics(keyword: string, sessionId?: string) {
+    const data = await this.topicModel.queryByKeyword(keyword, this.toDbSessionId(sessionId));
+
+    return data as unknown as Promise<ChatTopic[]>;
+  }
+
+  async getAllTopics() {
+    const data = await this.topicModel.queryAll();
+
+    return data as unknown as Promise<ChatTopic[]>;
+  }
+
+  async countTopics() {
+    return this.topicModel.count();
+  }
+
+  async updateTopic(id: string, data: Partial<ChatTopic>) {
+    return this.topicModel.update(id, data as any);
+  }
+
+  async removeTopic(id: string) {
+    return this.topicModel.delete(id);
+  }
+
+  async removeTopics(sessionId: string) {
+    return this.topicModel.batchDeleteBySessionId(this.toDbSessionId(sessionId));
+  }
+
+  async batchRemoveTopics(topics: string[]) {
+    return this.topicModel.batchDelete(topics);
+  }
+
+  async removeAllTopic() {
+    return this.topicModel.deleteAll();
+  }
+
+  private toDbSessionId(sessionId?: string | null) {
+    return sessionId === INBOX_SESSION_ID ? null : sessionId;
+  }
+}

package/src/services/user/{client.test.ts → _deprecated.test.ts}

@@ -4,9 +4,8 @@ import { Mock, beforeEach, describe, expect, it, vi } from 'vitest';
 import { UserModel } from '@/database/_deprecated/models/user';
 import { UserPreference } from '@/types/user';
 import { UserSettings } from '@/types/user/settings';
-import { AsyncLocalStorage } from '@/utils/localStorage';
 
-import { ClientService } from './client';
+import { ClientService } from './_deprecated';
 
 vi.mock('@/database/_deprecated/models/user', () => ({
   UserModel: {

package/src/services/user/index.ts

@@ -1,5 +1,11 @@
-import { ClientService } from './client';
+import { ClientService as DeprecatedService } from './_deprecated';
+import { ClientService } from './pglite';
 import { ServerService } from './server';
 
+const clientService =
+  process.env.NEXT_PUBLIC_CLIENT_DB === 'pglite' ? new ClientService() : new DeprecatedService();
+
 export const userService =
-  process.env.NEXT_PUBLIC_SERVICE_MODE === 'server' ? new ServerService() : new ClientService();
+  process.env.NEXT_PUBLIC_SERVICE_MODE === 'server' ? new ServerService() : clientService;
+
+export const userClientService = clientService;

package/src/services/user/pglite.test.ts

@@ -0,0 +1,98 @@
+import { eq } from 'drizzle-orm';
+import { DeepPartial } from 'utility-types';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { clientDB, initializeDB } from '@/database/client/db';
+import { userSettings, users } from '@/database/schemas';
+import { UserPreference } from '@/types/user';
+import { UserSettings } from '@/types/user/settings';
+
+import { ClientService } from './pglite';
+
+const mockUser = {
+  avatar: 'avatar.png',
+  settings: { themeMode: 'light' } as unknown as UserSettings,
+  uuid: 'user-id',
+};
+
+const mockPreference = {
+  useCmdEnterToSend: true,
+} as UserPreference;
+const clientService = new ClientService(mockUser.uuid);
+
+beforeEach(async () => {
+  vi.clearAllMocks();
+
+  await initializeDB();
+  await clientDB.delete(users);
+
+  await clientDB.insert(users).values({ id: mockUser.uuid, avatar: 'avatar.png' });
+  await clientDB
+    .insert(userSettings)
+    .values({ id: mockUser.uuid, general: { themeMode: 'light' } });
+});
+
+describe('ClientService', () => {
+  it('should get user state correctly', async () => {
+    const spyOn = vi
+      .spyOn(clientService['preferenceStorage'], 'getFromLocalStorage')
+      .mockResolvedValue(mockPreference);
+
+    const userState = await clientService.getUserState();
+
+    expect(userState).toMatchObject({
+      avatar: mockUser.avatar,
+      isOnboard: true,
+      canEnablePWAGuide: false,
+      hasConversation: false,
+      canEnableTrace: false,
+      preference: mockPreference,
+      settings: { general: { themeMode: 'light' } },
+      userId: mockUser.uuid,
+    });
+    expect(spyOn).toHaveBeenCalledTimes(1);
+  });
+
+  it('should update user settings correctly', async () => {
+    const settingsPatch: DeepPartial<UserSettings> = { general: { themeMode: 'dark' } };
+
+    await clientService.updateUserSettings(settingsPatch);
+
+    const result = await clientDB.query.userSettings.findFirst({
+      where: eq(userSettings.id, mockUser.uuid),
+    });
+
+    expect(result).toMatchObject(settingsPatch);
+  });
+
+  it('should reset user settings correctly', async () => {
+    await clientService.resetUserSettings();
+
+    const result = await clientDB.query.userSettings.findFirst({
+      where: eq(userSettings.id, mockUser.uuid),
+    });
+
+    expect(result).toBeUndefined();
+  });
+
+  it('should update user avatar correctly', async () => {
+    const newAvatar = 'new-avatar.png';
+
+    await clientService.updateAvatar(newAvatar);
+  });
+
+  it('should update user preference correctly', async () => {
+    const newPreference = {
+      useCmdEnterToSend: false,
+    } as UserPreference;
+
+    const spyOn = vi
+      .spyOn(clientService['preferenceStorage'], 'saveToLocalStorage')
+      .mockResolvedValue(undefined);
+
+    await clientService.updatePreference(newPreference);
+
+    expect(spyOn).toHaveBeenCalledWith(newPreference);
+    expect(spyOn).toHaveBeenCalledTimes(1);
+  });
+});

package/src/services/user/pglite.ts

@@ -0,0 +1,92 @@
+import { DeepPartial } from 'utility-types';
+
+import { clientDB } from '@/database/client/db';
+import { users } from '@/database/schemas';
+import { MessageModel } from '@/database/server/models/message';
+import { SessionModel } from '@/database/server/models/session';
+import { UserModel } from '@/database/server/models/user';
+import { BaseClientService } from '@/services/baseClientService';
+import { UserGuide, UserInitializationState, UserPreference } from '@/types/user';
+import { UserSettings } from '@/types/user/settings';
+import { AsyncLocalStorage } from '@/utils/localStorage';
+
+import { IUserService } from './type';
+
+export class ClientService extends BaseClientService implements IUserService {
+  private preferenceStorage: AsyncLocalStorage<UserPreference>;
+
+  private get userModel(): UserModel {
+    return new UserModel(clientDB as any, this.userId);
+  }
+  private get messageModel(): MessageModel {
+    return new MessageModel(clientDB as any, this.userId);
+  }
+  private get sessionModel(): SessionModel {
+    return new SessionModel(clientDB as any, this.userId);
+  }
+
+  constructor(userId?: string) {
+    super(userId);
+    this.preferenceStorage = new AsyncLocalStorage('LOBE_PREFERENCE');
+  }
+
+  async getUserState(): Promise<UserInitializationState> {
+    // if user not exist in the db, create one to make sure the user exist
+    await this.makeSureUserExist();
+
+    const state = await this.userModel.getUserState((encryptKeyVaultsStr) =>
+      encryptKeyVaultsStr ? JSON.parse(encryptKeyVaultsStr) : {},
+    );
+
+    const user = await UserModel.findById(clientDB as any, this.userId);
+    const messageCount = await this.messageModel.count();
+    const sessionCount = await this.sessionModel.count();
+
+    return {
+      ...state,
+      avatar: user?.avatar as string,
+      canEnablePWAGuide: messageCount >= 4,
+      canEnableTrace: messageCount >= 4,
+      hasConversation: messageCount > 0 || sessionCount > 0,
+      isOnboard: true,
+      preference: await this.preferenceStorage.getFromLocalStorage(),
+    };
+  }
+
+  updateUserSettings = async (value: DeepPartial<UserSettings>) => {
+    const { keyVaults, ...res } = value;
+
+    return this.userModel.updateSetting({ ...res, keyVaults: JSON.stringify(keyVaults) });
+  };
+
+  resetUserSettings = async () => {
+    return this.userModel.deleteSetting();
+  };
+
+  async updateAvatar(avatar: string) {
+    await this.userModel.updateUser({ avatar });
+  }
+
+  async updatePreference(preference: Partial<UserPreference>) {
+    await this.preferenceStorage.saveToLocalStorage(preference);
+  }
+
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars,unused-imports/no-unused-vars
+  async updateGuide(guide: Partial<UserGuide>) {
+    throw new Error('Method not implemented.');
+  }
+
+  async makeSureUserExist() {
+    const existUsers = await clientDB.query.users.findMany();
+
+    let user: { id: string };
+    if (existUsers.length === 0) {
+      const result = await clientDB.insert(users).values({ id: this.userId }).returning();
+      user = result[0];
+    } else {
+      user = existUsers[0];
+    }
+
+    return user;
+  }
+}

package/src/store/chat/slices/builtinTool/action.test.ts

@@ -2,7 +2,7 @@ import { act, renderHook } from '@testing-library/react';
 import { describe, expect, it, vi } from 'vitest';
 
 import { fileService } from '@/services/file';
-import { ClientService } from '@/services/file/client';
+import { ClientService } from '@/services/file/_deprecated';
 import { messageService } from '@/services/message';
 import { imageGenerationService } from '@/services/textToImage';
 import { uploadService } from '@/services/upload';
@@ -47,10 +47,9 @@ describe('chatToolSlice', () => {
         url: '',
       });
       vi.spyOn(result.current, 'toggleDallEImageLoading');
-      vi.spyOn(ClientService.prototype, 'checkFileHash').mockImplementation(async () => ({
-        isExist: false,
-        metadata: {},
-      }));
+      vi.spyOn(ClientService.prototype, 'checkFileHash').mockImplementation(
+        async () => ({ isExist: false }) as any,
+      );
 
       await act(async () => {
         await result.current.generateImageFromPrompts(prompts, messageId);

package/src/store/global/actions/clientDb.ts

@@ -0,0 +1,51 @@
+import { SWRResponse } from 'swr';
+import type { StateCreator } from 'zustand/vanilla';
+
+import { useOnlyFetchOnceSWR } from '@/libs/swr';
+import type { GlobalStore } from '@/store/global';
+import { DatabaseLoadingState, OnStageChange } from '@/types/clientDB';
+
+type InitClientDBParams = { onStateChange: OnStageChange };
+/**
+ * 设置操作
+ */
+export interface GlobalClientDBAction {
+  initializeClientDB: (params?: InitClientDBParams) => Promise<void>;
+  markPgliteEnabled: () => void;
+  useInitClientDB: (params?: InitClientDBParams) => SWRResponse;
+}
+
+export const clientDBSlice: StateCreator<
+  GlobalStore,
+  [['zustand/devtools', never]],
+  [],
+  GlobalClientDBAction
+> = (set, get) => ({
+  initializeClientDB: async (params) => {
+    // if the db has started initialized or not error, just skip.
+    if (
+      get().initClientDBStage !== DatabaseLoadingState.Idle &&
+      get().initClientDBStage !== DatabaseLoadingState.Error
+    )
+      return;
+
+    const { initializeDB } = await import('@/database/client/db');
+    await initializeDB({
+      onError: (error) => {
+        set({ initClientDBError: error });
+      },
+      onProgress: (data) => {
+        set({ initClientDBProcess: data });
+      },
+      onStateChange: (state) => {
+        set({ initClientDBStage: state });
+        params?.onStateChange?.(state);
+      },
+    });
+  },
+  markPgliteEnabled: () => {
+    get().updateSystemStatus({ isEnablePglite: true });
+  },
+  useInitClientDB: (params) =>
+    useOnlyFetchOnceSWR('initClientDB', () => get().initializeClientDB(params)),
+});

package/src/store/global/initialState.ts

@@ -1,5 +1,6 @@
 import { AppRouterInstance } from 'next/dist/shared/lib/app-router-context.shared-runtime';
 
+import { DatabaseLoadingState } from '@/types/clientDB';
 import { SessionDefaultGroup } from '@/types/session';
 import { AsyncLocalStorage } from '@/utils/localStorage';
 
@@ -37,6 +38,10 @@ export interface SystemStatus {
   hidePWAInstaller?: boolean;
   hideThreadLimitAlert?: boolean;
   inputHeight: number;
+  /**
+   * 应用初始化时不启用 PGLite，只有当用户手动开启时才启用
+   */
+  isEnablePglite?: boolean;
   mobileShowPortal?: boolean;
   mobileShowTopic?: boolean;
   sessionsWidth: number;
@@ -50,6 +55,13 @@ export interface SystemStatus {
 
 export interface GlobalState {
   hasNewVersion?: boolean;
+  initClientDBError?: Error;
+  initClientDBProcess?: { costTime?: number; phase: 'wasm' | 'dependencies'; progress: number };
+  /**
+   * 客户端数据库初始化状态
+   * 启动时为 Idle，完成为 Ready，报错为 Error
+   */
+  initClientDBStage: DatabaseLoadingState;
   isMobile?: boolean;
   isStatusInit?: boolean;
   latestVersion?: string;
@@ -76,6 +88,7 @@ export const INITIAL_STATUS = {
 } satisfies SystemStatus;
 
 export const initialState: GlobalState = {
+  initClientDBStage: DatabaseLoadingState.Idle,
   isMobile: false,
   isStatusInit: false,
   sidebarKey: SidebarTabKey.Chat,

package/src/store/global/selectors.ts

@@ -1,4 +1,6 @@
+import { isServerMode, isUsePgliteDB } from '@/const/version';
 import { GlobalStore } from '@/store/global';
+import { DatabaseLoadingState } from '@/types/clientDB';
 
 import { INITIAL_STATUS } from './initialState';
 
@@ -22,17 +24,36 @@ const filePanelWidth = (s: GlobalStore) => s.status.filePanelWidth;
 const inputHeight = (s: GlobalStore) => s.status.inputHeight;
 const threadInputHeight = (s: GlobalStore) => s.status.threadInputHeight;
 
-const isPgliteNotEnabled = () => false;
+const isPgliteNotEnabled = (s: GlobalStore) =>
+  isUsePgliteDB && !isServerMode && s.isStatusInit && !s.status.isEnablePglite;
 
-const isPgliteNotInited = () => false;
+/**
+ * 当且仅当 client db 模式，且 pglite 未初始化完成时返回 true
+ */
+const isPgliteNotInited = (s: GlobalStore) =>
+  isUsePgliteDB &&
+  s.isStatusInit &&
+  s.status.isEnablePglite &&
+  s.initClientDBStage !== DatabaseLoadingState.Ready;
 
-const isPgliteInited = (): boolean => true;
+/**
+ * 当且仅当 client db 模式，且 pglite 初始化完成时返回 true
+ */
+const isPgliteInited = (s: GlobalStore): boolean =>
+  (s.isStatusInit &&
+    s.status.isEnablePglite &&
+    s.initClientDBStage === DatabaseLoadingState.Ready) ||
+  false;
+
+// 这个变量控制 clientdb 是否完成初始化，正常来说，只有 pgliteDB 模式下，才会存在变化，其他时候都是 true
+const isDBInited = (s: GlobalStore): boolean => (isUsePgliteDB ? isPgliteInited(s) : true);
 
 export const systemStatusSelectors = {
   filePanelWidth,
   hidePWAInstaller,
   inZenMode,
   inputHeight,
+  isDBInited,
   isPgliteInited,
   isPgliteNotEnabled,
   isPgliteNotInited,

package/src/store/global/store.ts

@@ -5,15 +5,17 @@ import { StateCreator } from 'zustand/vanilla';
 
 import { createDevtools } from '../middleware/createDevtools';
 import { type GlobalStoreAction, globalActionSlice } from './action';
+import { type GlobalClientDBAction, clientDBSlice } from './actions/clientDb';
 import { type GlobalState, initialState } from './initialState';
 
 //  ===============  聚合 createStoreFn ============ //
 
-export type GlobalStore = GlobalState & GlobalStoreAction;
+export type GlobalStore = GlobalState & GlobalStoreAction & GlobalClientDBAction;
 
 const createStore: StateCreator<GlobalStore, [['zustand/devtools', never]]> = (...parameters) => ({
   ...initialState,
   ...globalActionSlice(...parameters),
+  ...clientDBSlice(...parameters),
 });
 
 //  ===============  实装 useStore ============ //