@lobehub/chat 0.140.0 → 0.141.0

package/src/const/settings.ts

@@ -6,6 +6,7 @@ import {
   GlobalDefaultAgent,
   GlobalLLMConfig,
   GlobalSettings,
+  GlobalSyncSettings,
   GlobalTTSConfig,
 } from '@/types/settings';
 
@@ -121,9 +122,14 @@ export const DEFAULT_TOOL_CONFIG = {
   },
 };
 
+const DEFAULT_SYNC_CONFIG: GlobalSyncSettings = {
+  webrtc: { enabled: false },
+};
+
 export const DEFAULT_SETTINGS: GlobalSettings = {
   defaultAgent: DEFAULT_AGENT,
   languageModel: DEFAULT_LLM_CONFIG,
+  sync: DEFAULT_SYNC_CONFIG,
   tool: DEFAULT_TOOL_CONFIG,
   tts: DEFAULT_TTS_CONFIG,
   ...DEFAULT_BASE_SETTINGS,

package/src/database/core/__tests__/model.test.ts

@@ -37,7 +37,7 @@ describe('BaseModel', () => {
         content: 'Hello, World!',
       };
 
-      const result = await baseModel['_add'](validData);
+      const result = await baseModel['_addWithSync'](validData);
 
       expect(result).toHaveProperty('id');
       expect(console.error).not.toHaveBeenCalled();
@@ -49,7 +49,7 @@ describe('BaseModel', () => {
         content: 'Hello, World!',
       };
 
-      await expect(baseModel['_add'](invalidData)).rejects.toThrow(TypeError);
+      await expect(baseModel['_addWithSync'](invalidData)).rejects.toThrow(TypeError);
     });
   });
 });

package/src/database/core/db.ts

@@ -21,7 +21,7 @@ import {
 } from './schemas';
 import { DBModel, LOBE_CHAT_LOCAL_DB_NAME } from './types/db';
 
-interface LobeDBSchemaMap {
+export interface LobeDBSchemaMap {
   files: DB_File;
   messages: DB_Message;
   plugins: DB_Plugin;

package/src/database/core/index.ts

@@ -1,2 +1,3 @@
 export { LocalDBInstance } from './db';
 export * from './model';
+export { dataSync } from './sync';

package/src/database/core/model.ts

@@ -1,10 +1,11 @@
 import Dexie, { BulkError } from 'dexie';
 import { ZodObject } from 'zod';
 
-import { DBBaseFieldsSchema } from '@/database/core/types/db';
 import { nanoid } from '@/utils/uuid';
 
 import { LocalDB, LocalDBInstance, LocalDBSchema } from './db';
+import { dataSync } from './sync';
+import { DBBaseFieldsSchema } from './types/db';
 
 export class BaseModel<N extends keyof LocalDBSchema = any, T = LocalDBSchema[N]['table']> {
   protected readonly db: LocalDB;
@@ -21,10 +22,16 @@ export class BaseModel<N extends keyof LocalDBSchema = any, T = LocalDBSchema[N]
     return this.db[this._tableName] as Dexie.Table;
   }
 
+  get yMap() {
+    return dataSync.getYMap(this._tableName);
+  }
+
+  // **************** Create *************** //
+
   /**
    * create a new record
    */
-  protected async _add<T = LocalDBSchema[N]['model']>(
+  protected async _addWithSync<T = LocalDBSchema[N]['model']>(
     data: T,
     id: string | number = nanoid(),
     primaryKey: string = 'id',
@@ -51,6 +58,9 @@ export class BaseModel<N extends keyof LocalDBSchema = any, T = LocalDBSchema[N]
 
     const newId = await this.db[tableName].add(record);
 
+    // sync data to yjs data map
+    this.updateYMapItem(newId);
+
     return { id: newId };
   }
 
@@ -69,6 +79,7 @@ export class BaseModel<N extends keyof LocalDBSchema = any, T = LocalDBSchema[N]
        */
       createWithNewId?: boolean;
       idGenerator?: () => string;
+      withSync?: boolean;
     } = {},
   ): Promise<{
     added: number;
@@ -77,8 +88,8 @@ export class BaseModel<N extends keyof LocalDBSchema = any, T = LocalDBSchema[N]
     skips: string[];
     success: boolean;
   }> {
-    const { idGenerator = nanoid, createWithNewId = false } = options;
-    const validatedData = [];
+    const { idGenerator = nanoid, createWithNewId = false, withSync = true } = options;
+    const validatedData: any[] = [];
     const errors = [];
     const skips: string[] = [];
 
@@ -123,6 +134,15 @@ export class BaseModel<N extends keyof LocalDBSchema = any, T = LocalDBSchema[N]
     try {
       await this.table.bulkAdd(validatedData);
 
+      if (withSync) {
+        dataSync.transact(() => {
+          const pools = validatedData.map(async (item) => {
+            await this.updateYMapItem(item.id);
+          });
+          Promise.all(pools);
+        });
+      }
+
       return {
         added: validatedData.length,
         ids: validatedData.map((item) => item.id),
@@ -144,7 +164,36 @@ export class BaseModel<N extends keyof LocalDBSchema = any, T = LocalDBSchema[N]
     }
   }
 
-  protected async _update(id: string, data: Partial<T>) {
+  // **************** Delete *************** //
+
+  protected async _deleteWithSync(id: string) {
+    const result = await this.table.delete(id);
+    // sync delete data to yjs data map
+    this.yMap?.delete(id);
+    return result;
+  }
+
+  protected async _bulkDeleteWithSync(keys: string[]) {
+    await this.table.bulkDelete(keys);
+    // sync delete data to yjs data map
+
+    dataSync.transact(() => {
+      keys.forEach((id) => {
+        this.yMap?.delete(id);
+      });
+    });
+  }
+
+  protected async _clearWithSync() {
+    const result = await this.table.clear();
+    // sync clear data to yjs data map
+    this.yMap?.clear();
+    return result;
+  }
+
+  // **************** Update *************** //
+
+  protected async _updateWithSync(id: string, data: Partial<T>) {
     // we need to check whether the data is valid
     // pick data related schema from the full schema
     const keys = Object.keys(data);
@@ -162,6 +211,35 @@ export class BaseModel<N extends keyof LocalDBSchema = any, T = LocalDBSchema[N]
 
     const success = await this.table.update(id, { ...data, updatedAt: Date.now() });
 
+    // sync data to yjs data map
+    this.updateYMapItem(id);
+
     return { success };
   }
+
+  protected async _putWithSync(data: any, id: string) {
+    const result = await this.table.put(data, id);
+
+    // sync data to yjs data map
+    this.updateYMapItem(id);
+
+    return result;
+  }
+
+  protected async _bulkPutWithSync(items: T[]) {
+    await this.table.bulkPut(items);
+
+    await dataSync.transact(() => {
+      items.forEach((items) => {
+        this.updateYMapItem((items as any).id);
+      });
+    });
+  }
+
+  // **************** Helper *************** //
+
+  private updateYMapItem = async (id: string) => {
+    const newData = await this.table.get(id);
+    this.yMap?.set(id, newData);
+  };
 }

package/src/database/core/sync.ts

@@ -0,0 +1,328 @@
+import Debug from 'debug';
+import { throttle, uniqBy } from 'lodash-es';
+import type { WebrtcProvider } from 'y-webrtc';
+import type { Doc, Transaction } from 'yjs';
+
+import {
+  OnAwarenessChange,
+  OnSyncEvent,
+  OnSyncStatusChange,
+  PeerSyncStatus,
+  StartDataSyncParams,
+} from '@/types/sync';
+
+import { LobeDBSchemaMap, LocalDBInstance } from './db';
+
+const LOG_NAME_SPACE = 'DataSync';
+
+class DataSync {
+  private _ydoc: Doc | null = null;
+  private provider: WebrtcProvider | null = null;
+
+  private syncParams!: StartDataSyncParams;
+  private onAwarenessChange!: OnAwarenessChange;
+
+  private waitForConnecting: any;
+
+  logger = Debug(LOG_NAME_SPACE);
+
+  transact(fn: (transaction: Transaction) => unknown) {
+    this._ydoc?.transact(fn);
+  }
+
+  getYMap = (tableKey: keyof LobeDBSchemaMap) => {
+    return this._ydoc?.getMap(tableKey);
+  };
+
+  startDataSync = async (params: StartDataSyncParams) => {
+    this.syncParams = params;
+    this.onAwarenessChange = params.onAwarenessChange;
+
+    // 开发时由于存在 fast refresh 全局实例会缓存在运行时中
+    // 因此需要在每次重新连接时清理上一次的实例
+    if (window.__ONLY_USE_FOR_CLEANUP_IN_DEV) {
+      await this.cleanConnection(window.__ONLY_USE_FOR_CLEANUP_IN_DEV);
+    }
+
+    await this.connect(params);
+  };
+
+  connect = async (params: StartDataSyncParams) => {
+    const {
+      channel,
+      onSyncEvent,
+      onSyncStatusChange,
+      user,
+      onAwarenessChange,
+      signaling = 'wss://y-webrtc-signaling.lobehub.com',
+    } = params;
+    // ====== 1. init yjs doc ====== //
+
+    await this.initYDoc();
+
+    this.logger('[YJS] start to listen sync event...');
+    this.initYjsObserve(onSyncEvent, onSyncStatusChange);
+
+    // ====== 2. init webrtc provider ====== //
+    this.logger(`[WebRTC] init provider... room: ${channel.name}`);
+    const { WebrtcProvider } = await import('y-webrtc');
+
+    // clients connected to the same room-name share document updates
+    this.provider = new WebrtcProvider(channel.name, this._ydoc!, {
+      password: channel.password,
+      signaling: [signaling],
+    });
+
+    // when fast refresh in dev, the provider will be cached in window
+    // so we need to clean it in destory
+    if (process.env.NODE_ENV === 'development') {
+      window.__ONLY_USE_FOR_CLEANUP_IN_DEV = this.provider;
+    }
+
+    this.logger(`[WebRTC] provider init success`);
+
+    // ====== 3. check signaling server connection  ====== //
+
+    // 当本地设备正确连接到 WebRTC Provider 后，触发 status 事件
+    // 当开始连接，则开始监听事件
+    this.provider.on('status', async ({ connected }) => {
+      this.logger('[WebRTC] peer status:', connected);
+      if (connected) {
+        // this.initObserve(onSyncEvent, onSyncStatusChange);
+        onSyncStatusChange?.(PeerSyncStatus.Connecting);
+      }
+    });
+
+    // check the connection with signaling server
+    let connectionCheckCount = 0;
+
+    this.waitForConnecting = setInterval(() => {
+      const signalingConnection: IWebsocketClient = this.provider!.signalingConns[0];
+
+      if (signalingConnection.connected) {
+        onSyncStatusChange?.(PeerSyncStatus.Ready);
+        clearInterval(this.waitForConnecting);
+        return;
+      }
+
+      connectionCheckCount += 1;
+
+      // check for 5 times, or make it failed
+      if (connectionCheckCount > 5) {
+        onSyncStatusChange?.(PeerSyncStatus.Unconnected);
+        clearInterval(this.waitForConnecting);
+      }
+    }, 2000);
+
+    // ====== 4. handle data sync  ====== //
+
+    // 当各方的数据均完成同步后，YJS 对象之间的数据已经一致时，触发 synced 事件
+    this.provider.on('synced', async ({ synced }) => {
+      this.logger('[WebRTC] peer sync status:', synced);
+      if (synced) {
+        this.logger('[WebRTC] start to init yjs data...');
+        onSyncStatusChange?.(PeerSyncStatus.Syncing);
+        await this.initSync();
+        onSyncStatusChange?.(PeerSyncStatus.Synced);
+        this.logger('[WebRTC] yjs data init success');
+      } else {
+        this.logger('[WebRTC] data not sync, try to reconnect in 1s...');
+        // await this.reconnect(params);
+        setTimeout(() => {
+          onSyncStatusChange?.(PeerSyncStatus.Syncing);
+          this.reconnect(params);
+        }, 1000);
+      }
+    });
+
+    // ====== 5. handle awareness  ====== //
+
+    this.initAwareness({ onAwarenessChange, user });
+
+    return this.provider;
+  };
+
+  reconnect = async (params: StartDataSyncParams) => {
+    await this.cleanConnection(this.provider);
+
+    await this.connect(params);
+  };
+
+  async disconnect() {
+    await this.cleanConnection(this.provider);
+  }
+
+  private initYDoc = async () => {
+    if (typeof window === 'undefined') return;
+
+    this.logger('[YJS] init YDoc...');
+    const { Doc } = await import('yjs');
+    this._ydoc = new Doc();
+  };
+
+  private async cleanConnection(provider: WebrtcProvider | null) {
+    if (provider) {
+      this.logger(`[WebRTC] clean Connection...`);
+      this.logger(`[WebRTC] clean awareness...`);
+      provider.awareness.destroy();
+
+      this.logger(`[WebRTC] clean room...`);
+      provider.room?.disconnect();
+      provider.room?.destroy();
+
+      this.logger(`[WebRTC] clean provider...`);
+      provider.disconnect();
+      provider.destroy();
+
+      this.logger(`[WebRTC] clean yjs doc...`);
+      this._ydoc?.destroy();
+
+      this.logger(`[WebRTC] -------------------`);
+    }
+  }
+
+  private initSync = async () => {
+    await Promise.all(
+      ['sessions', 'sessionGroups', 'topics', 'messages', 'plugins'].map(async (tableKey) =>
+        this.loadDataFromDBtoYjs(tableKey as keyof LobeDBSchemaMap),
+      ),
+    );
+  };
+
+  private initYjsObserve = (onEvent: OnSyncEvent, onSyncStatusChange: OnSyncStatusChange) => {
+    ['sessions', 'sessionGroups', 'topics', 'messages', 'plugins'].forEach((tableKey) => {
+      // listen yjs change
+      this.observeYMapChange(tableKey as keyof LobeDBSchemaMap, onEvent, onSyncStatusChange);
+    });
+  };
+
+  private observeYMapChange = (
+    tableKey: keyof LobeDBSchemaMap,
+    onEvent: OnSyncEvent,
+    onSyncStatusChange: OnSyncStatusChange,
+  ) => {
+    const table = LocalDBInstance[tableKey];
+    const yItemMap = this.getYMap(tableKey);
+    const updateSyncEvent = throttle(onEvent, 1000);
+
+    // 定义一个变量来保存定时器的ID
+    // eslint-disable-next-line no-undef
+    let debounceTimer: NodeJS.Timeout;
+
+    yItemMap?.observe(async (event) => {
+      // abort local change
+      if (event.transaction.local) return;
+
+      // 每次有变更时，都先清除之前的定时器（如果有的话），然后设置新的定时器
+      clearTimeout(debounceTimer);
+
+      onSyncStatusChange(PeerSyncStatus.Syncing);
+
+      this.logger(`[YJS] observe ${tableKey} changes:`, event.keysChanged.size);
+      const pools = Array.from(event.keys).map(async ([id, payload]) => {
+        const item: any = yItemMap.get(id);
+
+        switch (payload.action) {
+          case 'add':
+          case 'update': {
+            await table.put(item, id);
+
+            break;
+          }
+
+          case 'delete': {
+            await table.delete(id);
+            break;
+          }
+        }
+      });
+
+      await Promise.all(pools);
+
+      updateSyncEvent(tableKey);
+
+      // 设置定时器，2000ms 后更新状态为'synced'
+      debounceTimer = setTimeout(() => {
+        onSyncStatusChange(PeerSyncStatus.Synced);
+      }, 2000);
+    });
+  };
+
+  private loadDataFromDBtoYjs = async (tableKey: keyof LobeDBSchemaMap) => {
+    const table = LocalDBInstance[tableKey];
+    const items = await table.toArray();
+    const yItemMap = this.getYMap(tableKey);
+
+    // 定义每批次最多包含的数据条数
+    const batchSize = 50;
+
+    // 计算总批次数
+    const totalBatches = Math.ceil(items.length / batchSize);
+
+    for (let i = 0; i < totalBatches; i++) {
+      // 计算当前批次的起始和结束索引
+      const start = i * batchSize;
+      const end = start + batchSize;
+
+      // 获取当前批次的数据
+      const batchItems = items.slice(start, end);
+
+      // 将当前批次的数据推送到 Yjs 中
+      this._ydoc?.transact(() => {
+        batchItems.forEach((item) => {
+          yItemMap!.set(item.id, item);
+        });
+      });
+    }
+
+    this.logger('[DB]:', tableKey, yItemMap?.size);
+  };
+
+  private initAwareness = ({ user }: Pick<StartDataSyncParams, 'user' | 'onAwarenessChange'>) => {
+    if (!this.provider) return;
+
+    const awareness = this.provider.awareness;
+
+    awareness.setLocalState({ clientID: awareness.clientID, user });
+    this.onAwarenessChange?.([{ ...user, clientID: awareness.clientID, current: true }]);
+
+    awareness.on('change', () => this.syncAwarenessToUI());
+  };
+
+  private syncAwarenessToUI = async () => {
+    const awareness = this.provider?.awareness;
+
+    if (!awareness) return;
+
+    const state = Array.from(awareness.getStates().values()).map((s) => ({
+      ...s.user,
+      clientID: s.clientID,
+      current: s.clientID === awareness.clientID,
+    }));
+
+    this.onAwarenessChange?.(uniqBy(state, 'id'));
+  };
+}
+
+export const dataSync = new DataSync();
+
+interface IWebsocketClient {
+  binaryType: 'arraybuffer' | 'blob' | null;
+  connect(): void;
+  connected: boolean;
+  connecting: boolean;
+  destroy(): void;
+  disconnect(): void;
+  lastMessageReceived: number;
+  send(message: any): void;
+  shouldConnect: boolean;
+  unsuccessfulReconnects: number;
+  url: string;
+  ws: WebSocket;
+}
+
+declare global {
+  interface Window {
+    __ONLY_USE_FOR_CLEANUP_IN_DEV?: WebrtcProvider | null;
+  }
+}

package/src/database/models/__tests__/message.test.ts

@@ -1,6 +1,5 @@
 import { afterEach, beforeEach, describe, expect, it } from 'vitest';
 
-import { DB_Message } from '@/database/schemas/message';
 import { ChatMessage } from '@/types/message';
 
 import { CreateMessageParams, MessageModel } from '../message';

package/src/database/models/__tests__/plugin.test.ts

@@ -59,11 +59,14 @@ describe('PluginModel', () => {
   describe('update', () => {
     it('should update a plugin', async () => {
       await PluginModel.create(pluginData);
-      const updatedPluginData: DB_Plugin = { ...pluginData, type: 'customPlugin' };
+      const updatedPluginData: DB_Plugin = {
+        ...pluginData,
+        type: 'customPlugin',
+      };
       await PluginModel.update(pluginData.identifier, updatedPluginData);
       const plugins = await PluginModel.getList();
       expect(plugins).toHaveLength(1);
-      expect(plugins[0]).toEqual(updatedPluginData);
+      expect(plugins[0]).toEqual({ ...updatedPluginData, updatedAt: expect.any(Number) });
     });
   });
 

package/src/database/models/file.ts

@@ -11,7 +11,7 @@ class _FileModel extends BaseModel<'files'> {
   async create(file: DB_File) {
     const id = nanoid();
 
-    return this._add(file, `file-${id}`);
+    return this._addWithSync(file, `file-${id}`);
   }
 
   async findById(id: string) {

package/src/database/models/message.ts

@@ -107,13 +107,13 @@ class _MessageModel extends BaseModel {
 
     const messageData: DB_Message = this.mapChatMessageToDBMessage(data as ChatMessage);
 
-    return this._add(messageData, id);
+    return this._addWithSync(messageData, id);
   }
 
   async batchCreate(messages: ChatMessage[]) {
     const data: DB_Message[] = messages.map((m) => this.mapChatMessageToDBMessage(m));
 
-    return this._batchAdd(data);
+    return this._batchAdd(data, { withSync: true });
   }
 
   async duplicateMessages(messages: ChatMessage[]): Promise<ChatMessage[]> {
@@ -123,36 +123,14 @@ class _MessageModel extends BaseModel {
     return duplicatedMessages;
   }
 
-  async createDuplicateMessages(messages: ChatMessage[]): Promise<ChatMessage[]> {
-    // 创建一个映射来存储原始消息ID和复制消息ID之间的关系
-    const idMapping = new Map<string, string>();
-
-    // 首先复制所有消息，并为每个复制的消息生成新的ID
-    const duplicatedMessages = messages.map((originalMessage) => {
-      const newId = nanoid();
-      idMapping.set(originalMessage.id, newId);
-
-      return { ...originalMessage, id: newId };
-    });
-
-    // 更新 parentId 为复制后的新ID
-    for (const duplicatedMessage of duplicatedMessages) {
-      if (duplicatedMessage.parentId && idMapping.has(duplicatedMessage.parentId)) {
-        duplicatedMessage.parentId = idMapping.get(duplicatedMessage.parentId);
-      }
-    }
-
-    return duplicatedMessages;
-  }
-
   // **************** Delete *************** //
 
   async delete(id: string) {
-    return this.table.delete(id);
+    return super._deleteWithSync(id);
   }
 
   async clearTable() {
-    return this.table.clear();
+    return this._clearWithSync();
   }
 
   /**
@@ -178,13 +156,32 @@ class _MessageModel extends BaseModel {
     const messageIds = await query.primaryKeys();
 
     // Use the bulkDelete method to delete all selected messages in bulk
-    return this.table.bulkDelete(messageIds);
+    return this._bulkDeleteWithSync(messageIds);
+  }
+
+  async batchDeleteBySessionId(sessionId: string): Promise<void> {
+    // If topicId is specified, use both assistantId and topicId as the filter criteria in the query.
+    // Otherwise, filter by assistantId and require that topicId is undefined.
+    const messageIds = await this.table.where('sessionId').equals(sessionId).primaryKeys();
+
+    // Use the bulkDelete method to delete all selected messages in bulk
+    return this._bulkDeleteWithSync(messageIds);
+  }
+
+  /**
+   * Delete all messages associated with the topicId
+   * @param topicId
+   */
+  async batchDeleteByTopicId(topicId: string): Promise<void> {
+    const messageIds = await this.table.where('topicId').equals(topicId).primaryKeys();
+
+    return this._bulkDeleteWithSync(messageIds);
   }
 
   // **************** Update *************** //
 
   async update(id: string, data: DeepPartial<DB_Message>) {
-    return this._update(id, data);
+    return super._updateWithSync(id, data);
   }
 
   async updatePluginState(id: string, key: string, value: any) {
@@ -202,7 +199,7 @@ class _MessageModel extends BaseModel {
    */
   async batchUpdate(messageIds: string[], updateFields: Partial<DB_Message>): Promise<number> {
     // Retrieve the messages by their IDs
-    const messagesToUpdate = await this.table.where(':id').anyOf(messageIds).toArray();
+    const messagesToUpdate = await this.table.where('id').anyOf(messageIds).toArray();
 
     // Update the specified fields of each message
     const updatedMessages = messagesToUpdate.map((message) => ({
@@ -211,13 +208,35 @@ class _MessageModel extends BaseModel {
     }));
 
     // Use the bulkPut method to update the messages in bulk
-    await this.table.bulkPut(updatedMessages);
+    await this._bulkPutWithSync(updatedMessages);
 
     return updatedMessages.length;
   }
 
   // **************** Helper *************** //
 
+  private async createDuplicateMessages(messages: ChatMessage[]): Promise<ChatMessage[]> {
+    // 创建一个映射来存储原始消息ID和复制消息ID之间的关系
+    const idMapping = new Map<string, string>();
+
+    // 首先复制所有消息，并为每个复制的消息生成新的ID
+    const duplicatedMessages = messages.map((originalMessage) => {
+      const newId = nanoid();
+      idMapping.set(originalMessage.id, newId);
+
+      return { ...originalMessage, id: newId };
+    });
+
+    // 更新 parentId 为复制后的新ID
+    for (const duplicatedMessage of duplicatedMessages) {
+      if (duplicatedMessage.parentId && idMapping.has(duplicatedMessage.parentId)) {
+        duplicatedMessage.parentId = idMapping.get(duplicatedMessage.parentId);
+      }
+    }
+
+    return duplicatedMessages;
+  }
+
   private mapChatMessageToDBMessage(message: ChatMessage): DB_Message {
     const { extra, ...messageData } = message;